Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`. Args: seed (`int`): The seed to set. Returns: `None`
(seed: int)
| 54 | |
| 55 | |
| 56 | def set_seed(seed: int): |
| 57 | """ |
| 58 | Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`. |
| 59 | |
| 60 | Args: |
| 61 | seed (`int`): The seed to set. |
| 62 | |
| 63 | Returns: |
| 64 | `None` |
| 65 | """ |
| 66 | random.seed(seed) |
| 67 | np.random.seed(seed) |
| 68 | torch.manual_seed(seed) |
| 69 | if is_torch_npu_available(): |
| 70 | torch.npu.manual_seed_all(seed) |
| 71 | else: |
| 72 | torch.cuda.manual_seed_all(seed) |
| 73 | # ^^ safe to call this function even if cuda is not available |
| 74 | |
| 75 | |
| 76 | def compute_snr(noise_scheduler, timesteps): |
searching dependent graphs…