Saving and Loading in Python

A trained PyTorch model exists as a graph of Python objects sitting in RAM (or GPU memory) — the moment the process that holds it exits, that state is gone. Forty-eight hours of GPU training and a 99% accuracy score are worth nothing if you close the interpreter without explicitly writing the learned parameters to disk first.

This is why serialization is not an optional step tacked onto the end of a training script; it's the boundary between 'a model that exists' and 'a model you can actually use.' A production pipeline saves checkpoints during training (so a crash doesn't erase hours of progress) and saves a final artifact after training completes, so the model can be loaded again later — in a notebook, on a web server, or on a completely different machine.

PyTorch's answer to this problem is torch.save(), and the object it's almost always used to save is the model's state_dict — the learned weights and biases — rather than the model instance itself. The next sections unpack exactly what that dictionary contains and why it's the recommended unit of serialization.

Every nn.Module in PyTorch exposes its learned parameters through .state_dict(), which returns an OrderedDict mapping each layer's name (as defined in the model's __init__) to the tensor holding its current weights or biases — for example "fc1.weight" or "conv1.bias".

Crucially, the state_dict contains only numbers: the parameter tensors as they currently stand after training. It does not contain the Python class definition, the forward() method's logic, or any information about how those layers are connected — that architecture only exists in the code that defines the MyNetwork class itself.

Calling model.state_dict().keys() is a fast way to sanity-check what a model actually learned to store — every layer that has trainable parameters shows up as a key, in the exact order it was registered on the module.

It's worth being precise about what model.state_dict() is not, because the wrong mental model here causes real bugs later. It is not a copy of your training data, not the Python source file, and not a snapshot of the optimizer's internal state (momentum buffers, learning-rate schedule) — that's a separate dictionary, optimizer.state_dict(), that you'd save alongside it only if you intend to resume training later.

What it is: a dictionary containing every learned weight and bias tensor for every layer that has trainable parameters, keyed by layer name. If your network has a self.fc1 = nn.Linear(784, 128), the returned dict will have entries like "fc1.weight" (shape [128, 784]) and "fc1.bias" (shape [128]).

This distinction matters the moment you go to save a model: saving just model.state_dict() is enough to run inference later, but if your pipeline needs to resume interrupted training, you need a checkpoint that bundles the model's state_dict together with the optimizer's state_dict and the current epoch number.

torch.save() serializes a Python object to disk using a pickle-based binary format, with special handling so tensor storage is written efficiently rather than pickled naively. Calling torch.save(model.state_dict(), "my_awesome_model.pth") writes out just the weights-and-biases dictionary from the previous section, not the model instance itself.

The .pt and .pth extensions are pure convention — PyTorch doesn't inspect the filename to decide how to serialize anything, and either extension (or none at all) works identically. What actually matters is what you pass as the first argument: a state_dict produces a small, portable file containing only tensors; passing the model object itself produces a much more fragile artifact, which the next section explains.

Because the file is just a serialized dictionary of tensors, it can be inspected, copied, or version-controlled like any other binary artifact — there's no hidden 'live' Python process bundled inside it.

PyTorch technically lets you call torch.save(model, "model.pth") to pickle the entire model object, architecture and all. It's tempting because loading it back looks simpler — torch.load("model.pth") alone hands you a ready-to-use model, no class instantiation required. That convenience is also exactly why it's discouraged for anything beyond a quick local experiment.

Pickling the full object ties the saved file to the *exact* class definition, module path, and even directory structure that existed at save time. Rename the class, move the file that defines it, or refactor your project layout, and loading the pickled model breaks with an import error — even though the underlying weights are perfectly fine.

The recommended pattern avoids all of this: save only model.state_dict(), and keep the class definition itself as ordinary version-controlled source code. Loading then becomes a two-step, code-independent process — instantiate the architecture from source, then inject the saved weights — which the next two sections walk through.

Reloading a model saved via state_dict is a deliberate two-step process. First, you instantiate the model class exactly as it was defined during training — loaded_model = MyNetwork() — which creates a fresh module with randomly initialized weights and the correct layer structure, but none of the learned knowledge yet.

Second, you call loaded_model.load_state_dict(torch.load("my_awesome_model.pth")). torch.load() deserializes the saved file back into an OrderedDict of tensors, and load_state_dict() walks that dictionary, matching each key to the corresponding parameter on the freshly created model and copying the saved values in place.

Once that call succeeds, loaded_model is functionally identical to the model at the moment it was saved — same weights, same biases, same learned behavior — even though it was built from a brand-new instance of the class.

The saved .pth file contains raw numbers keyed by layer name — "fc1.weight", "fc1.bias", and so on — but nothing in that dictionary describes how those layers are wired together, what activation functions sit between them, or what forward() should do with an input tensor. That logic lives entirely in the MyNetwork class definition, which is ordinary Python code, not something torch.save() ever captured.

So instantiating MyNetwork() again isn't a formality — it's what recreates the actual computation graph and gives load_state_dict() a set of correctly-shaped parameter slots to copy values into. Without that empty architecture to load into, the saved tensors would just be a dictionary of numbers with no model to attach them to.

This is also why load_state_dict() fails loudly (RuntimeError: Missing key(s) or Unexpected key(s)) if the class definition has changed since saving — the keys in the checkpoint no longer line up with the parameters the fresh instance actually has.

The two-step save/load pattern solves the code-portability problem, but there's a second, separate portability issue lurking in the saved tensors themselves: the hardware they were trained on. Every tensor in a state_dict remembers which device it lived on — CPU or a specific GPU index — at the moment it was saved.

That device information travels with the file. If you don't account for it explicitly when loading on different hardware, torch.load() will try to recreate tensors on the exact same device they came from, which fails outright if that device doesn't exist on the machine doing the loading.

The next two sections walk through exactly when this bites you and the one-argument fix that makes a checkpoint hardware-agnostic.

Picture the concrete failure: you train on an NVIDIA GPU, so every tensor in the saved state_dict is tagged cuda:0. You email the .pth file to a friend running a laptop with no NVIDIA GPU at all. They run loaded_model.load_state_dict(torch.load("my_awesome_model.pth")) and instead of a working model, they get a RuntimeError — PyTorch tried to allocate those tensors on a CUDA device that simply doesn't exist on their machine.

This isn't a bug in their setup; it's torch.load() doing exactly what it's documented to do by default — restoring each tensor onto the device it was originally saved from. The file itself is perfectly valid; the mismatch is purely about where the bytes get materialized.

This exact scenario is extremely common in practice: models are frequently trained on GPU clusters and then deployed to CPU-only inference servers, or shared with collaborators who don't have matching hardware — which is precisely the problem the next section's fix addresses.

The fix is a single keyword argument on torch.load(): map_location. Passing torch.load("my_awesome_model.pth", map_location=torch.device('cpu')) tells PyTorch to remap every tensor in the checkpoint onto CPU memory during deserialization, regardless of which device they were originally saved from — no CUDA device required.

This is the standard, hardware-agnostic way to load a checkpoint of unknown origin. A common production pattern goes one step further and makes it dynamic: device = torch.device('cuda' if torch.cuda.is_available() else 'cpu'), then torch.load(path, map_location=device) — the exact same loading code then works correctly whether it runs on a GPU server or a CPU-only laptop.

After loading with map_location, the tensors sit on the target device, but the model itself should still be moved there explicitly with loaded_model.to(device) before running inference, so the model and its inputs agree on where they live.

Put together, the full deployment workflow covered in this lesson is: train the model, serialize only its state_dict with torch.save() rather than the whole object, and on the receiving end instantiate the same architecture from source before calling load_state_dict() to restore the weights — with map_location handling any mismatch between the training and deployment hardware.

This pattern is exactly what production ML deployment relies on: a checkpoint file that's small, portable across machines, and decoupled from any specific pickle-fragile object graph. Whether the model ends up behind a REST API, inside a batch inference job, or loaded into a Jupyter notebook for evaluation, the loading code looks the same.

From here, the next step is combining this with more advanced checkpointing — bundling the model's state_dict together with the optimizer's state and the current epoch, so training can be paused and resumed rather than just re-run from scratch.

You trained a model for 48 hours on a GPU. It has 99% accuracy. If you close Python right now, all that training is deleted instantly.

In PyTorch, the matrix of learned weights and biases is stored in a Python dictionary called the state_dict.

What exactly does model.state_dict() contain in PyTorch?

• →The source code of the Python file.
• →It is a Python dictionary containing all the learned weights and biases for every layer in the neural network.
• →A copy of the training data.

To save the model securely, you use torch.save(), passing it the state_dict and a filename (usually ending in .pt or .pth).

What is the PyTorch best practice for saving a trained model to the hard drive?

• →Use Pandas to save it as a CSV.
• →Save ONLY the state_dict (the weights) using torch.save(), rather than trying to save the entire Python class.
• →Take a screenshot of the terminal.

To load the model on a web server tomorrow, you must first instantiate the EMPTY class architecture, and then load the weights into it.

Why must you instantiate the MyNetwork() class AGAIN before you can load your saved weights?

• →Because Python automatically deletes old variables.
• →Because the saved .pth file only contains the raw numbers (weights), not the architecture logic. PyTorch needs the class framework to map the numbers to the right layers.
• →To bypass NVIDIA licensing.

Now, prepare yourself. We are about to enter the ADA Defense Protocol. Ensure you understand device mapping during loading.

You trained the model on an NVIDIA GPU. You send the .pth file to a friend on a MacBook without an NVIDIA GPU. If they run torch.load(), it crashes.

ADA DEFENSE: How do you safely load a PyTorch model that was trained on a GPU onto a machine that only has a CPU?

• →You cannot. The friend must buy an NVIDIA GPU.
• →Pass map_location=torch.device('cpu') into the torch.load() function to force the tensors into CPU RAM.
• →Rename the .pth file to .cpu.

Threat neutralized. Deployment protocols secured. Proceeding to Advanced Architectures.

Build a Real State Dictionary. Finish build_state_dict(): map each layer's name to its learned weights.