How do you save and load PyTorch models correctly, including full training checkpoints?

PyTorch supports saving either the full model object or just its weights (state_dict). Saving only the state_dict is the recommended approach because it decouples weights from the Python class definition. A full training checkpoint includes the optimizer state too, so training can resume exactly where it left off.

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 only
torch.save(model.state_dict(), "weights.pt")

model_new = nn.Linear(10, 5)        # must define the SAME architecture first
model_new.load_state_dict(torch.load("weights.pt"))
model_new.eval()                     # always call before inference

# ── NOT recommended: save the entire model object
# Fragile — breaks if the class definition moves or changes
torch.save(model, "full_model.pt")
loaded_model = torch.load("full_model.pt", weights_only=False)

# ── Full training checkpoint — for resuming training
def save_checkpoint(path, epoch, model, optimizer, best_val_loss):
    torch.save({
        "epoch":            epoch,
        "model_state":      model.state_dict(),
        "optimizer_state":  optimizer.state_dict(),  # Adam momentum buffers etc.
        "best_val_loss":    best_val_loss,
    }, path)

def load_checkpoint(path, model, optimizer):
    ckpt = torch.load(path, map_location="cpu")  # always load to CPU first
    model.load_state_dict(ckpt["model_state"])
    optimizer.load_state_dict(ckpt["optimizer_state"])
    return ckpt["epoch"], ckpt["best_val_loss"]

save_checkpoint("ckpt.pt", epoch=5, model=model, optimizer=optimizer, best_val_loss=0.42)
epoch, best_loss = load_checkpoint("ckpt.pt", model_new, optimizer)

# ── Loading on a different device than it was saved
model.load_state_dict(
    torch.load("weights.pt", map_location="cpu")  # avoid GPU OOM if GPU unavailable
)
model = model.to("cuda")  # then move to the desired device

Invest now in Acorns!!! 🚀
Join Acorns and get your $5 bonus!

Acorns is a micro-investing app that automatically invests your "spare change" from daily purchases into diversified, expert-built portfolios of ETFs. It is designed for beginners, allowing you to start investing with as little as $5. The service automates saving and investing. Disclosure: I may receive a referral bonus.

Invest now!!! Get Free equity stock (US, UK only)!

Use Robinhood app to invest in stocks. It is safe and secure. Use the Referral link to claim your free stock when you sign up!.

The Robinhood app makes it easy to trade stocks, crypto and more.

Webull! Receive free stock by signing up using the link: Webull signup.