PyTorch Save and Load Models

Saving and loading PyTorch models lets you keep trained weights, resume interrupted training and deploy a model elsewhere. The recommended approach is to save the model's state_dict, a dictionary mapping each layer name to its weight tensor, rather than pickling the whole model object.

What is a state_dict?

import torch
from torch import nn

model = nn.Sequential(nn.Linear(4, 2), nn.ReLU(), nn.Linear(2, 1))
for k, v in model.state_dict().items():
    print(k, tuple(v.shape))

0.weight (2, 4)
0.bias (2,)
2.weight (1, 2)
2.bias (1,)

Saving and Loading Weights

torch.save(model.state_dict(), "model.pt")

new_model = nn.Sequential(nn.Linear(4, 2), nn.ReLU(), nn.Linear(2, 1))
new_model.load_state_dict(torch.load("model.pt", weights_only=True))
new_model.eval()
print("loaded")

loaded

You must first build a model with the same architecture, then load the weights into it. Call eval() before inference so dropout and batch norm behave correctly.

Since PyTorch 2.6, torch.load defaults to weights_only=True, which refuses to unpickle arbitrary Python objects. That protects you from malicious files; only turn it off for files you fully trust.

Checkpoints for Resuming Training

opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
epoch = 7
torch.save({
    "epoch": epoch,
    "model": model.state_dict(),
    "optimizer": opt.state_dict(),
}, "checkpoint.pt")

ckpt = torch.load("checkpoint.pt", weights_only=True)
model.load_state_dict(ckpt["model"])
opt.load_state_dict(ckpt["optimizer"])
start_epoch = ckpt["epoch"] + 1
print("resume from epoch", start_epoch)

resume from epoch 8

Saving the optimizer state matters for Adam-style optimizers, which keep running averages for every parameter.

Loading GPU Weights on a CPU

state = torch.load("model.pt", map_location="cpu", weights_only=True)
new_model.load_state_dict(state)

<All keys matched successfully>

Other Formats

• safetensors: a safe, fast format used across Hugging Face; save with safetensors.torch.save_file.
• torch.export: captures a model graph for deployment and ahead-of-time compilation.
• ONNX: export with torch.onnx.export to run in ONNX Runtime and other engines.
• Avoid torch.save(model) of the whole object: it breaks when your code files move or change.

Frequently asked questions

How do I save a PyTorch model?

Use torch.save(model.state_dict(), "model.pt"). To load, build the same architecture and call model.load_state_dict(torch.load("model.pt", weights_only=True)).

What is the difference between .pt and .pth files?

Nothing technical; both are just conventions for files written by torch.save. .pt is the more common choice today.

Why does torch.load fail with weights_only errors?

Since PyTorch 2.6, weights_only=True is the default and blocks arbitrary pickled objects. Save plain state_dicts, or pass weights_only=False only for files you trust.