PyTorch · Lesson 7 · Capstone

Ship It & Read It

Save a trained model and load it back to identical predictions; learn the checkpoint pattern for resuming training; then read a complete real-world training script top to bottom — and realize you now understand every line.

🎯 Mission: read & write real PyTorch
The one idea

You don't save "the model" — you save its state_dict: a plain dictionary of its learned tensors. To restore, you rebuild the same architecture in code and pour the saved numbers back in. That single habit — code defines the shape, the file holds the weights — is how every PyTorch model is shipped, resumed, and shared.

01 Why this is the last piece

Training is the expensive part. Once a model is good, you persist it so you can load it instantly for predictions, hand it to someone else, or resume a long training run after a crash. That's "shipping" — and it's the last mechanic standing between you and the full workflow.

This lesson also doubles as your graduation check: in §06 you'll read a real training script and find nothing mysterious in it. That's the mission — read & write real PyTorch — met.

02 Save the state_dict, not the model

The recommended way to save is the model's state_dict — the dictionary of its parameter tensors you met in Lesson 4:

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

print(list(model.state_dict().keys()))
# ['0.weight', '0.bias', '2.weight', '2.bias']   ← Sequential names by position
# (a custom module with self.fc would show 'fc.weight', 'fc.bias')
💡 Why not just torch.save(model)? Saving the whole object pickles your Python class and file paths with it — brittle, and it breaks if your code moves or changes. The state_dict is just the numbers: small, portable, and decoupled from your code layout. Our 82-param net's file is ~2.5 KB. Convention: name the file .pt or .pth.

Source: PyTorch — Saving & Loading Models.

03 Load: rebuild the architecture, then pour in the weights

Loading is two steps, and the order matters: construct the same model in code, then load the saved state_dict into it. Finish with eval() before predicting.

fresh = build()                              # SAME architecture (random weights for now)
fresh.load_state_dict(torch.load("model.pt"))  # overwrite with the saved numbers
fresh.eval()                                 # inference mode before predicting
output before saving: [[-0.2084,  0.2171]]
fresh (random) output: [[-0.2518, -0.1944]]   # different — not loaded yet
loaded output:         [[-0.2084,  0.2171]]   # identical after load_state_dict
identical to original? True
trained model .state_dict() model.pt weights on disk build() — same architecture load_state_dict eval() predict save load code defines the shape · the file holds the weights
The round trip. Because the file is only numbers, you must recreate the architecture in code before loading.
⚠ The mismatch error load_state_dict fails if the architecture doesn't match the saved keys/shapes ("Missing/Unexpected key(s)"). The model you build to load into must be the same as the one you saved. And call eval() before inference (Lesson 6) — forgetting it is a silent accuracy bug on models with dropout/batch-norm.

04 Checkpoints — to resume training

To resume a long run (not just predict), you need more than weights: the optimizer's state and which epoch you were on. The idiom is to save a plain dict:

torch.save({
    "epoch": epoch,
    "model_state": model.state_dict(),
    "optim_state": opt.state_dict(),   # Adam's momentum etc. — needed to resume cleanly
}, "checkpoint.pt")

ckpt = torch.load("checkpoint.pt")
model.load_state_dict(ckpt["model_state"])
opt.load_state_dict(ckpt["optim_state"])
start_epoch = ckpt["epoch"]          # pick up where you left off
💡 Weights = inference, checkpoint = resume Ship just the state_dict when all you need is to run the model. Save a full checkpoint (model + optimizer + epoch, and anything else) when you might need to continue training. You'll see both in real repos.

05 Loading across devices

A model saved from the GPU can be loaded on a CPU-only machine with map_location:

state = torch.load("model.pt", map_location="cpu")   # or map_location=device
⚠ A modern default to know: weights_only Recent PyTorch defaults torch.load(..., weights_only=True), which only un-pickles tensors — a security guard against loading untrusted files that could run arbitrary code. For plain state_dicts (what you should be saving) this just works; you only relax it for files you trust that contain non-tensor objects.
Run the whole save/load demo: ~/projects/learn/public/courses/pytorch/practice/.venv/bin/python ~/projects/learn/public/courses/pytorch/practice/lesson07_save_load.py

06 Capstone: read a real training script

Here is a complete, idiomatic PyTorch training program — the kind you'll meet in the wild. Read it slowly. Every construct has a number pointing to the lesson that taught it. Nothing here should be a mystery anymore.

import torch, torch.nn as nn
from torch.utils.data import DataLoader, random_split, TensorDataset

device = "mps" if torch.backends.mps.is_available() else "cpu"   ⟶ L1 (device)

ds = TensorDataset(X, y)                                          ⟶ L5 (Dataset)
train_ds, val_ds = random_split(ds, [640, 160])                    ⟶ L5 (split)
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)   ⟶ L5 (batch/shuffle)

model = nn.Sequential(nn.Linear(2,16), nn.ReLU(), nn.Linear(16,2)).to(device)  ⟶ L4/L6 (model)
loss_fn = nn.CrossEntropyLoss()                                  ⟶ L6 (loss)
opt = torch.optim.Adam(model.parameters(), lr=1e-2)              ⟶ L4 (optimizer)

for epoch in range(31):                                          ⟶ L3 (epochs)
    model.train()                                                ⟶ L6 (train mode)
    for xb, yb in train_loader:                                  ⟶ L5 (batches)
        xb, yb = xb.to(device), yb.to(device)                    ⟶ L1/L6 (to device)
        opt.zero_grad()                                          ⟶ L2/L4 (reset grads)
        loss = loss_fn(model(xb), yb)                            ⟶ L3 (forward + loss)
        loss.backward()                                          ⟶ L2 (gradients)
        opt.step()                                               ⟶ L4 (update)

torch.save(model.state_dict(), "model.pt")                       ⟶ L7 (ship it)
💡 A reading checklist for any PyTorch script Find these six things and you understand the whole file: (1) the data (Dataset/DataLoader, batch size), (2) the model (layers), (3) the loss, (4) the optimizer + learning rate, (5) the loop (the five moves), (6) the device + save/eval. Everything else is detail layered on this skeleton.

07 You've hit the mission — where next

🎓 Mission check You can now read a tensor by its shape/dtype/device, explain what autograd computes, write the training loop from memory, build an idiomatic nn.Module with a real loss and optimizer, feed it with a DataLoader, train on your GPU, and ship the result. That's reading & writing real PyTorch.

Natural next directions when you want them — each reuses everything here, only swapping the model or data:

CNNs for images

nn.Conv2d + torchvision datasets. The loop is unchanged; the model learns spatial features. (The 60-Minute Blitz ends here.)

Transformers & LLMs

The architecture behind modern AI. Same tensors/autograd/loop — a likely next mission for you.

Read real repos

Apply the §06 checklist to a project on GitHub. Reading code is now your fastest way to level up.

💡 Wisdom comes from the community When you hit something the lessons didn't cover, the PyTorch Forums are official and well-moderated (core devs answer). For beginner Q&A and code review, r/pytorch. Bring a tiny reproducible snippet — half the time, writing it reveals the answer. (See RESOURCES.md.)

08 Check yourself

What's the recommended thing to save?
Before calling load_state_dict, what must you do?
You want to resume training later. What do you save?
A model trained on MPS, loading on a CPU-only machine. What helps?
Reading an unfamiliar PyTorch script, which six things orient you fastest?

09 Flashcards

Q: What do you save to ship a model, and why that?
A: Its state_dict (a dict of parameter tensors). It's portable and decoupled from your code, unlike pickling the whole object.
Q: The two-step load procedure?
A: (1) Build the same architecture in code; (2) model.load_state_dict(torch.load(path)). Then model.eval() before inference.
Q: state_dict vs checkpoint — when each?
A: state_dict when you only need to run the model. Checkpoint (model + optimizer state + epoch) when you may resume training.
Q: How do you load a GPU-saved model on CPU?
A: torch.load(path, map_location="cpu") (or map_location=device).
Q: What is weights_only=True protecting against?
A: Un-pickling untrusted files that could execute arbitrary code. It's the modern default and works fine for plain state_dicts.
Q: The six-point checklist for reading any training script?
A: data (Dataset/DataLoader), model (layers), loss, optimizer + lr, the loop (five moves), device + save/eval.
👩‍🏫 I'm your teacher — and that's the core course

You've gone from "what's a tensor" to reading and writing a full training script — the mission. From here, learning is best driven by your goals: bring me a repo you want to understand, a model you want to build, or say the word "transformers" and we'll start the next mission. Want a quick scenario quiz to pressure-test all seven lessons? Just ask.