03 · Loading for Inference 🚀

4 min read

Training a model and running it are two different jobs — often on two different machines. Loading for inference is about doing the second job safely and portably.

1. The 'Why'

The whole point of saving a model is to use it later, and "later" often means somewhere very different from where it was trained. You might train on a beefy CUDA server and then load the model on a colleague's Mac (MPS) or on a modest CPU-only production box. A model's saved tensors remember which device they lived on, so naively loading a GPU-trained checkpoint on a CPU-only machine throws an error. Inference also has its own correctness requirement: the model must be in evaluation mode and gradient tracking must be off, or your predictions will be both wrong (Dropout still firing) and needlessly slow and memory-hungry (a graph being built for a backward pass that never comes).

This sub-module ties together threads from across the course into the clean, repeatable recipe for inference. The map_location argument handles device portability, remapping saved tensors onto whatever hardware you're loading onto. The eval() + torch.no_grad() combination — which you've seen building since Module 02 — is mandatory here. And there's a modern safety consideration: torch.load can execute arbitrary code when unpickling untrusted files, so PyTorch now defaults to a safer weights_only mode you should understand. Get this recipe right and your model runs correctly anywhere, from a laptop to a server.

2. Core Concepts

map_location. Tells torch.load where to place the loaded tensors, regardless of where they were saved. torch.load(path, map_location="cpu") loads a GPU-trained model onto a CPU. Common values: "cpu", a specific device, or a device you computed at runtime.

The inference recipe. Re-create the model → load_state_dict (with map_location as needed) → model.to(device) → model.eval() → run inside torch.no_grad(). Every step earns its place.

eval() + no_grad() are both required. eval() fixes layer behavior (Dropout/BatchNorm); no_grad() disables gradient tracking. Inference needs both — one without the other is a bug.

weights_only=True (safety). Recent PyTorch defaults torch.load to weights_only=True, which refuses to execute arbitrary pickled code and only restores tensors/state_dicts. Keep it on for anything you didn't create yourself; that's exactly why saving a bare state_dict (sub-module 01) is preferred.

3. Code in Action

The full, portable inference recipe

import torch
import torch.nn as nn

# Pick whatever device THIS machine has
device = "cuda" if torch.cuda.is_available() else "cpu"

# 1. Re-create the architecture
model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))

# 2. Load weights, remapping tensors onto this machine's device
state = torch.load("model_weights.pt", map_location=device)  # device-portable
model.load_state_dict(state)

# 3. Move model to the device and switch to inference mode
model.to(device)
model.eval()                      # Dropout/BatchNorm -> inference behavior

# 4. Predict without building a graph
with torch.no_grad():
    x = torch.randn(5, 4).to(device)   # new, unseen inputs on the same device
    logits = model(x)
    preds = logits.argmax(dim=1)       # predicted class per sample
    print(preds)

Loading a GPU-trained model onto a CPU

import torch

# On a CPU-only box, this would ERROR without map_location:
state = torch.load("gpu_trained.pt", map_location="cpu")  # remap to CPU
# model.load_state_dict(state)   # now succeeds anywhere

Turning logits into probabilities for reporting

import torch
import torch.nn.functional as F

with torch.no_grad():
    logits = model(torch.randn(1, 4).to(device))
    probs = F.softmax(logits, dim=1)   # apply softmax ONLY for human-readable output
    print(probs)                       # e.g. tensor([[0.12, 0.80, 0.08]])

4. Common Pitfalls

Pitfall 1 — Forgetting map_location across devices. Loading a CUDA checkpoint on a machine without a GPU raises "Attempting to deserialize object on a CUDA device..." Pass map_location="cpu" (or the target device) to load anywhere.

Pitfall 2 — Skipping eval() and/or no_grad(). Without eval(), Dropout randomizes your predictions and BatchNorm uses batch stats — results wobble and worsen. Without no_grad(), you waste memory building a useless graph. Inference needs both, every time.

Pitfall 3 — Loading untrusted files with code execution enabled. torch.load(..., weights_only=False) on a file from someone else can run arbitrary code. Prefer bare state_dict files and keep weights_only=True (the modern default) unless you fully trust the source.

5. Further Reading & Watch List

⬅️ Prev: 02 · Checkpoints & Resuming · ➡️ Next: 04 · Best Model & Early Stopping