Args: Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`. seed (`int`): The seed to set.
(seed: int)
| 35 | |
| 36 | |
| 37 | def set_seed(seed: int): |
| 38 | """ |
| 39 | Args: |
| 40 | Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`. |
| 41 | seed (`int`): The seed to set. |
| 42 | """ |
| 43 | random.seed(seed) |
| 44 | np.random.seed(seed) |
| 45 | torch.manual_seed(seed) |
| 46 | if is_torch_npu_available(): |
| 47 | torch.npu.manual_seed_all(seed) |
| 48 | else: |
| 49 | torch.cuda.manual_seed_all(seed) |
| 50 | # ^^ safe to call this function even if cuda is not available |
| 51 | |
| 52 | |
| 53 | def compute_snr(noise_scheduler, timesteps): |