How do you save and load PyTorch models correctly, and what is included in a proper checkpoint?

PyTorch provides two main ways to persist a model: saving the full model object (convenient but fragile to class definition changes) or saving only the state dictionary (recommended for production and reproducibility). The state dict is a Python OrderedDict mapping layer names to their parameter tensors — it contains everything needed to recreate the model's learned state.

A proper training checkpoint includes more than just model weights — it must also save the optimizer state (which contains momentum buffers and adaptive learning rate accumulators in Adam), the current epoch and step, the best validation metric, and the random number generator state, so that training can be resumed exactly where it left off without any change in behaviour.

import torch
import torch.nn as nn

model = nn.Linear(10, 5)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# ── Recommended: save/load state_dict ──
torch.save(model.state_dict(), 'model_weights.pt')

model_new = nn.Linear(10, 5)               # same architecture
model_new.load_state_dict(torch.load('model_weights.pt'))
model_new.eval()                            # ALWAYS call eval() for inference

# ── Full training checkpoint ──
def save_checkpoint(path, epoch, model, optimizer, best_val_loss):
    torch.save({
        'epoch':         epoch,
        'model_state':   model.state_dict(),
        'optimizer_state': optimizer.state_dict(),
        'best_val_loss': best_val_loss,
        'rng_state':     torch.get_rng_state(),
    }, path)

def load_checkpoint(path, model, optimizer):
    ckpt = torch.load(path, map_location='cpu')
    model.load_state_dict(ckpt['model_state'])
    optimizer.load_state_dict(ckpt['optimizer_state'])
    return ckpt['epoch'], ckpt['best_val_loss']

# Loading on different device: always load to CPU first,
# then move to device (avoids GPU OOM if original GPU is unavailable)
model.load_state_dict(
    torch.load('model_weights.pt', map_location='cpu')
)
model = model.to('cuda')