How to save and restore a PyTorch training checkpoint

Long PyTorch training runs need more than final model weights when a job stops before the next epoch completes. A training checkpoint captures the model tensors, optimizer buffers, epoch position, and recent loss so training can continue from the saved state instead of restarting from random initialization.

PyTorch stores model and optimizer state through state_dict dictionaries, and torch.save serializes those dictionaries into one checkpoint file. Restoring requires new model and optimizer objects with the same structure before load_state_dict applies the saved tensors.

Use weights_only=True when loading a checkpoint that contains tensors, primitive values, and dictionaries from a trusted training run. Set map_location to the device that should receive the tensors during restore, then switch the restored model back to training mode before the next optimizer step.

1. Create a smoke script that trains once, saves a checkpoint, restores new objects, and resumes training.

   checkpoint_resume.py

   import torch
   from torch import nn
    
    
   torch.manual_seed(7)
   checkpoint_path = "training-checkpoint.tar"
    
    
   def build_model():
       return nn.Sequential(
           nn.Linear(3, 4),
           nn.ReLU(),
           nn.Linear(4, 1),
       )
    
    
   x = torch.tensor(
       [
           [0.1, 0.2, 0.3],
           [0.4, 0.5, 0.6],
       ],
       dtype=torch.float32,
   )
   y = torch.tensor(
       [
           [0.6],
           [1.5],
       ],
       dtype=torch.float32,
   )
   loss_fn = nn.MSELoss()
    
    
   def train_step(model, optimizer):
       model.train()
       optimizer.zero_grad()
       loss = loss_fn(model(x), y)
       loss.backward()
       optimizer.step()
       return loss.item()
    
    
   model = build_model()
   optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    
   loss_before_save = train_step(model, optimizer)
   torch.save(
       {
           "epoch": 1,
           "model_state_dict": model.state_dict(),
           "optimizer_state_dict": optimizer.state_dict(),
           "loss": loss_before_save,
       },
       checkpoint_path,
   )
   print(f"saved epoch=1 loss={loss_before_save:.4f}")
    
   restored_model = build_model()
   restored_optimizer = torch.optim.Adam(restored_model.parameters(), lr=0.01)
   checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    
   restored_model.load_state_dict(checkpoint["model_state_dict"])
   restored_optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
   restored_model.train()
    
   state_match = all(
       torch.equal(model.state_dict()[name], restored_model.state_dict()[name])
       for name in model.state_dict()
   )
   optimizer_state_entries = len(restored_optimizer.state_dict()["state"])
   next_epoch = checkpoint["epoch"] + 1
   resumed_loss = train_step(restored_model, restored_optimizer)
    
   print(f"loaded keys={', '.join(checkpoint.keys())}")
   print(f"state_match={state_match}")
   print(f"optimizer_state_entries={optimizer_state_entries}")
   print(f"resume_epoch={next_epoch} previous_loss={checkpoint['loss']:.4f}")
   print(f"resumed_loss={resumed_loss:.4f}")

   Build the restored model and optimizer with the same architecture and optimizer class that created the checkpoint. load_state_dict loads values into existing objects; it does not recreate the model class.

2. Run the checkpoint smoke script.

   $ python checkpoint_resume.py
   saved epoch=1 loss=0.7571
   loaded keys=epoch, model_state_dict, optimizer_state_dict, loss
   state_match=True
   optimizer_state_entries=4
   resume_epoch=2 previous_loss=0.7571
   resumed_loss=0.7287

   state_match=True confirms that the model tensors loaded into the new model object. optimizer_state_entries=4 confirms that the Adam buffers were restored before the resumed step.

3. Save the same checkpoint fields after a completed training epoch in the real training loop.

   torch.save(
       {
           "epoch": epoch,
           "model_state_dict": model.state_dict(),
           "optimizer_state_dict": optimizer.state_dict(),
           "loss": float(loss.item()),
       },
       "training-checkpoint.tar",
   )

   Store enough metadata to choose the next epoch and compare resumed training logs with the interrupted run.

4. Restore the checkpoint before the next training loop starts.

   model = build_model()
   optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    
   checkpoint = torch.load(
       "training-checkpoint.tar",
       map_location="cpu",
       weights_only=True,
   )
   model.load_state_dict(checkpoint["model_state_dict"])
   optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
   start_epoch = checkpoint["epoch"] + 1
   model.train()

   Load only checkpoint files from a trusted source. torch.load deserializes saved data; weights_only=True restricts loading to tensors, primitive values, dictionaries, and allowlisted safe globals.

5. Remove the smoke-test files when the checkpoint pattern has been verified.

   $ rm checkpoint_resume.py training-checkpoint.tar