llm_lab/utils/seed.py

"""Seed utility for reproducibility."""
import torch


def set_seed(seed: int = 42):
    """Set seed for reproducibility."""
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)