Saving, Loading & Running Inference

Why saving models correctly matters#

Training can take hours or days. You need to save the result, resume if training stops, and load the model later for predictions in an app or API. PyTorch has a clear, recommended way to do each.

Save and load with state_dict#

A model's state_dict is a dictionary of its parameter tensors (weights and biases). Saving the state_dict — not the whole model object — is the recommended approach because it's portable and doesn't depend on pickling your Python class.

import torch
from torch import nn

def make_model():
    return nn.Sequential(
        nn.Linear(2, 32), nn.ReLU(),
        nn.Linear(32, 32), nn.ReLU(),
        nn.Linear(32, 1),
    )

model = make_model()
# ... train the model ...

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

To load, recreate the same architecture, then load the weights into it:

restored = make_model()
restored.load_state_dict(torch.load("circle_model.pt", weights_only=True))
restored.eval()

weights_only=True makes torch.load refuse to run arbitrary Python code hidden inside a file — always use it, and never load model files from sources you don't trust.

Running inference#

points = torch.tensor([[0.0, 0.0], [1.5, 1.5]])

with torch.inference_mode():
    probabilities = torch.sigmoid(restored(points))

print((probabilities > 0.5).flatten().tolist())

[True, False]

The point (0, 0) is inside the circle and (1.5, 1.5) is outside — the restored model works.

Inference checklist:

1. model.eval() — puts dropout and batch-norm layers into inference behaviour.
2. torch.inference_mode() — disables gradient tracking for speed and lower memory.
3. Same preprocessing as training — identical normalisation, image size or tokenizer.
4. Same device for model and inputs.

Checkpoints: resume training later#

To continue training after a crash or a session timeout, save more than the weights — save the optimizer state (Adam keeps running averages) and the epoch:

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

Resume:

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

Save a checkpoint every epoch (or every N steps) and keep the one with the best validation score.

Loading across devices#

A model trained on a GPU can be loaded on a CPU-only machine with map_location:

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

Serving the model in an API#

A minimal pattern for a web API (e.g. FastAPI) is:

• Load the model once at startup, call model.eval().
• For each request: validate input → convert to a tensor → run under inference_mode() → return JSON.
• Batch multiple requests together when traffic is high.

For production at scale, models are often exported to ONNX (torch.onnx.export) or served with dedicated inference servers; open-source LLMs typically use specialised LLM serving engines.

Speeding up inference#

• GPU and half precision (model.half() or bfloat16) for large models.
• torch.compile(model) to optimise execution on supported hardware.
• Quantization (8-bit or 4-bit weights) to shrink memory and speed up CPU inference.

Common mistakes#

• Saving the whole model with torch.save(model). It breaks when you rename or move your class; save the state_dict.
• Forgetting model.eval(). Predictions change randomly because dropout is still active.
• Loading untrusted .pt files without weights_only=True. Pickled files can execute code.
• Different preprocessing at inference time. The model silently performs worse.

Key takeaways#

• Save and load state_dicts; recreate the architecture before loading.
• Use weights_only=True and map_location when loading.
• Checkpoints include optimizer state and epoch so training can resume.
• For inference: eval(), inference_mode() and identical preprocessing.