File size: 491 Bytes
8e04e6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 | import lightning as L
from lightning.pytorch.callbacks import Callback
class SeedCallback(Callback):
def on_save_checkpoint(self, trainer, pl_module, checkpoint):
checkpoint["global_step"] = trainer.global_step
def on_load_checkpoint(self, trainer, pl_module, checkpoint):
global_step = checkpoint["global_step"]
seed = global_step + trainer.global_rank
L.seed_everything(seed)
print(f"Seeding rank {trainer.global_rank} with seed {seed}")
|