Saving and loading with state_dict

Save the weights, not the object

A module's state_dict() is a dictionary mapping names such as fc1.weight to tensors: every parameter plus buffers like BatchNorm's running averages. The recommended way to save a model is torch.save(model.state_dict(), "model.pt"). To load it, first construct the model class in code, then call model.load_state_dict(torch.load("model.pt")).

You can torch.save(model) the whole object, but that pickles a reference to your class and module path; renaming a file or class later breaks loading. Saving only the state dict keeps weights independent of code layout.

To resume training, save a checkpoint dictionary: the model's state dict, the optimizer's state dict (Adam's running averages are part of training state), the scheduler's, and the epoch number. Pass map_location="cpu" (or your device) to torch.load so weights saved on a GPU load on a machine without one. After loading for inference, call model.eval(). load_state_dict is strict by default and reports missing or unexpected keys, which usually means the architecture changed.

From trained model to service

A trained model is saved as a state_dict checkpoint, loaded and verified on the target machine, optionally exported to a portable format such as torch.export or ONNX, and finally served behind an application.

A resumable checkpoint

SmallCNN, model, optimizer and epoch come from the training loop. weights_only=True restricts unpickling to tensors and plain containers, which is all a checkpoint like this needs.

import torch

checkpoint = {
    "epoch": epoch,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}
torch.save(checkpoint, "checkpoint.pt")

# Later, possibly on another machine:
model = SmallCNN()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
ckpt = torch.load("checkpoint.pt", map_location="cpu", weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
start_epoch = ckpt["epoch"] + 1

torch.load runs pickle: treat files as code

torch.load uses Python's pickle format, and unpickling an untrusted file can execute arbitrary code. Recent PyTorch versions default to weights_only=True, which only allows tensors and simple types; never switch it to False for files you did not create. For sharing weights publicly, the safetensors format stores tensors with no executable content.

Quick check: What is the recommended way to save a trained model for later use?

• `torch.save(model)` to pickle the whole object
• `torch.save(model.state_dict(), path)`, then rebuild the class and call `load_state_dict`
• Print the weights and copy them into a text file
• Save only the optimizer's state dict