githubmirror/ColossalAI

colossalai/testing/random.py

19 строк542 B

Hongxin Liu

[misc] update pre-commit and run all files (#4752)

19 сен 2023, 09:20

19 сен 2023, 09:20079bf3c

import random

import numpy as np
import torch


def seed_all(seed, cuda_deterministic=False):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    if cuda_deterministic:  # slower, more reproducible
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    else:
        torch.backends.cudnn.deterministic = False
        torch.backends.cudnn.benchmark = True