INNER CODE UNIT · TypeScript
load_checkpoint
theaniketgiri/create-llm · src/python-callback-templates.ts:191
def load_checkpoint(self, checkpoint_path: str, trainer: Any):
"""Load a checkpoint"""
checkpoint = torch.load(checkpoint_path)
trainer.model.load_state_dict(checkpoint['model_state_dict'])
trainer.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
if 'scheduler_state_dict' in checkpoint and hasattr(trainer, 'scheduler'):
trainer.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
trainer.global_step = checkpoint['step']
trainer.epoch = checkpoint.get('epoch', 0)
if self.verbose:
print(f"Loaded checkpoint from: {checkpoint_path}")
print(f"Resuming from step {trainer.global_step}")
return checkpoint