Background
Back to Articles/PyTorch 101: Zero to Hero
Sections
IntroductionModule 01 — Tensors 🧊01 · Creating Tensors 🧊02 · Indexing & Reshaping 🔪03 · Tensor Math ➗04 · Device Placement 🖥️⚡Module 02 — Autograd ⚙️01 · The Computational Graph 🕸️02 · Backward Pass & Gradients ⬅️03 · Turning Autograd Off 🛑04 · Gradient Gotchas 🪤Module 03 — Neural Networks01 · The `nn.Module` Basics 🧱02 · Common Layers 🧩03 · Building a Network 🏗️04 · Inspecting Models 🔍Module 04 — Data Handling 🗂️01 · Dataset Basics 📇02 · The DataLoader 🚚03 · Transforms 🎨04 · Splits & Built-in Datasets ✂️Module 05 — The Training Loop 🔁01 · Loss Functions 🎯02 · Optimizers ⚙️03 · The Training Loop 🔁04 · Evaluation & Metrics 📊Module 06 — Saving & Loading 💾01 · `state_dict` Basics 💾02 · Checkpoints & Resuming ⏸️03 · Loading for Inference 🚀04 · Best Model & Early Stopping 🏅Module 07 — Computer Vision 👁️01 · Convolutions 🔲02 · CNN Architecture 🏛️03 · Transfer Learning 🔁04 · Image Classification Project 🧪Module 08 — NLP & Transformers 💬01 · Text Data & Tokenization 🔤02 · Embeddings 🧭03 · Recurrent Layers & LSTMs 🔄04 · Intro to Transformers ⚡Module 09 — Ecosystem: PyTorch Lightning ⚡01 · Why Lightning? 🤔02 · The LightningModule 🧩03 · The Trainer 🎛️04 · DataModules & Callbacks 🧰Module 10 — Deployment & Optimization 🚀01 · Exporting Models 📦02 · `torch.compile` ⚡03 · Inference Optimization 🪶04 · Serving & Next Steps 🎓

01 · `state_dict` Basics 💾

4 min read

There's a right way and a fragile way to save a PyTorch model. Learn the right way once, and your saved files will survive code changes and version bumps.

1. The 'Why'

When training finishes, everything your model has learned lives in its parameters — the weight and bias tensors inside each layer. Persisting a model therefore comes down to one question: how do you write those numbers to disk and read them back reliably? PyTorch gives you two options, and choosing correctly matters more than beginners expect. The tempting option is to pickle the entire model object in one line. It works today, on your machine, with your exact code — but it secretly welds the saved file to your source code's file paths and class definitions, so a harmless refactor later can make the file un-loadable.

The robust option, and the one the PyTorch team recommends, is to save only the state_dict — the plain dictionary of parameter names to tensors that you met in Module 03. Because it stores just the values, decoupled from your code's structure, a state_dict file stays loadable even as you reorganize modules or upgrade libraries. The small tradeoff is that you must re-create the model architecture in code before loading the weights into it — but that's a feature, not a bug: it keeps the "what the model is" (your code) separate from "what the model learned" (the saved values). This one habit prevents a whole category of frustrating, hard-to-debug loading failures.

2. Core Concepts

torch.save and torch.load. The general-purpose serialization pair. torch.save(obj, path) writes any picklable object (typically a state_dict); torch.load(path) reads it back.

Save the state_dict, not the model. torch.save(model.state_dict(), "model.pt") stores only the learned values. This is the recommended, portable approach.

Loading requires re-creating the model first. Because the state_dict has no architecture, you instantiate the model class, then call model.load_state_dict(torch.load(path)) to pour the values in.

load_state_dict is strict by default. The keys in the file must match the model's keys exactly, or it raises an error — a helpful safeguard that catches architecture mismatches. Pass strict=False to allow partial loads (useful in transfer learning, Module 07).

File extensions. .pt and .pth are conventional; they're ordinary files, the extension is just custom.

3. Code in Action

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))

# --- SAVE: store only the learned values ---
torch.save(model.state_dict(), "model_weights.pt")

# --- LOAD: re-create the architecture, then load values into it ---
loaded = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))  # same structure!
loaded.load_state_dict(torch.load("model_weights.pt"))
loaded.eval()                         # switch to inference mode after loading

Inspecting what's inside

import torch

sd = torch.load("model_weights.pt")   # a plain OrderedDict
for name, tensor in sd.items():
    print(name, tuple(tensor.shape))
# 0.weight (8, 4) / 0.bias (8,) / 2.weight (3, 8) / 2.bias (3,)

The fragile alternative (know it, avoid it)

import torch

# Saves the WHOLE object via pickle — ties the file to your code layout
torch.save(model, "whole_model.pt")           # ⚠️ fragile across refactors
reloaded = torch.load("whole_model.pt", weights_only=False)  # needs the class importable
# Prefer the state_dict approach above unless you have a specific reason.

4. Common Pitfalls

Pitfall 1 — Saving the whole model and getting burned later. torch.save(model, ...) pickles class references and paths. Move or rename the file that defines your model class and the load fails. Save the state_dict instead; it doesn't care where your code lives.

Pitfall 2 — Architecture mismatch on load. load_state_dict demands the model you're loading into has the exact same layers as the one you saved from. A different hidden size or layer order throws a keys-mismatch error. Keep your model definition in sync (or version it alongside the weights).

Pitfall 3 — Forgetting model.eval() after loading. A freshly loaded model defaults to training mode, so Dropout/BatchNorm misbehave during inference. Always call model.eval() after loading for prediction (covered fully in sub-module 03).

5. Further Reading & Watch List

➡️ Next: 02 · Checkpoints & Resuming