Saving and loading checkpoints

Long training runs crash and models must be shipped; checkpoints make both possible.

The recommended way to persist a model is its state_dict: an ordered dict mapping parameter and buffer names to tensors. Save it rather than the whole model object, because pickling the object ties the file to your exact class and module layout.

To resume training you need more than weights: save a checkpoint dict with the model, optimizer (Adam's moment estimates), scheduler, epoch and anything else you track.

Loading:

Rebuild the model with the same code, then model.load_state_dict(sd). It reports missing or unexpected keys; strict=False tolerates them, for example when loading a backbone into a model with a new head.
torch.load(path, map_location="cpu") loads GPU-saved tensors on a CPU-only machine.
Since PyTorch 2.6, torch.load defaults to weights_only=True, which only unpickles tensors and basic containers. That protects you from malicious pickle files; keep it on and avoid saving arbitrary objects.

Call model.eval() after loading for inference. Models wrapped by torch.compile or DistributedDataParallel prefix their keys (_orig_mod., module.), so save the underlying module's state_dict.

For sharing models outside Python, look at torch.export or the safetensors format used by the Hugging Face ecosystem.

Example

import torch
from torch import nn

model = nn.Linear(3, 1)
opt = torch.optim.Adam(model.parameters())
print(list(model.state_dict().keys()))     # ['weight', 'bias']

torch.save({"model": model.state_dict(),
            "opt": opt.state_dict(),
            "epoch": 5}, "ckpt.pt")

ckpt = torch.load("ckpt.pt", map_location="cpu", weights_only=True)
model2 = nn.Linear(3, 1)                   # rebuild the architecture in code
model2.load_state_dict(ckpt["model"])
opt2 = torch.optim.Adam(model2.parameters())
opt2.load_state_dict(ckpt["opt"])

print(ckpt["epoch"])                       # 5
print(torch.equal(model.weight, model2.weight))   # True

Remember

Save state_dict()s, not whole model objects
A resumable checkpoint includes optimizer state and the epoch
torch.load uses weights_only=True by default; use map_location across devices

Watch out for

Saving only the model weights and then resuming with a fresh Adam optimizer, which causes a loss spike.

Check yourself

Why is saving model.state_dict() preferred over torch.save(model)?

Official docs: pytorch.org/tutorials/beginner/saving_loading_models.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks