Introduction

In PyTorch, the standard way to save a trained model is to save its state_dict, not the entire model object. The state_dict contains the learned parameters, gives you a more stable loading workflow, and avoids tying the checkpoint too tightly to the exact Python module structure at save time.

Save the model weights

The most common pattern is:

1import torch
2import torch.nn as nn
3
4class Net(nn.Module):
5    def __init__(self):
6        super().__init__()
7        self.fc = nn.Linear(4, 2)
8
9    def forward(self, x):
10        return self.fc(x)
11
12model = Net()
13
14torch.save(model.state_dict(), "model_weights.pth")

This saves only the trained parameters, not the Python class definition itself.

Load the weights later

To load that checkpoint, recreate the model class first and then load the saved state dictionary.

1import torch
2
3model = Net()
4state_dict = torch.load("model_weights.pth", map_location="cpu")
5model.load_state_dict(state_dict)
6model.eval()

model.eval() matters for inference because layers such as dropout and batch normalization behave differently in training and evaluation modes.

Why state_dict is preferred

Saving the full model object is possible:

The reason many teams standardize on state_dict is not just style. It reduces coupling between the checkpoint file and the exact Python import path of the model class at the moment the file was created.

torch.save(model, "full_model.pth")

But that approach is more fragile. It relies on Python pickling and can break more easily if:

• the model class moves to a different module
• the code layout changes
• the environment differs between save and load time

That is why state_dict is the default recommendation for most projects.

Save a training checkpoint

If you want to resume training instead of only running inference, save more than just model weights.

1import torch
2import torch.optim as optim
3
4optimizer = optim.Adam(model.parameters(), lr=1e-3)
5epoch = 12
6loss_value = 0.034
7
8torch.save({
9    "epoch": epoch,
10    "model_state_dict": model.state_dict(),
11    "optimizer_state_dict": optimizer.state_dict(),
12    "loss": loss_value,
13}, "checkpoint.pth")

This is the normal pattern for training checkpoints.

Resume training from a checkpoint

1checkpoint = torch.load("checkpoint.pth", map_location="cpu")
2
3model = Net()
4optimizer = optim.Adam(model.parameters(), lr=1e-3)
5
6model.load_state_dict(checkpoint["model_state_dict"])
7optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
8start_epoch = checkpoint["epoch"] + 1
9loss_value = checkpoint["loss"]
10
11model.train()

Now the training state is restored rather than only the learned weights.

CPU and GPU considerations

If a model was saved on one device and loaded on another, use map_location appropriately.

This is especially important when a model trained on CUDA hardware is later restored in a CPU-only deployment or test environment.

state_dict = torch.load("model_weights.pth", map_location="cpu")

This is especially important when loading a GPU-trained checkpoint on a CPU-only environment.

Naming and organization

A practical project usually keeps different checkpoint types separate:

That small bit of operational discipline makes deployment, rollback, and experiment tracking much easier once a project grows beyond a single notebook or script.

• latest checkpoint
• best validation checkpoint
• final inference weights

That makes deployment and recovery easier. Good checkpoint naming matters more than many teams expect.

Common Pitfalls

A common mistake is saving the full model object and later discovering that loading breaks after refactoring the codebase.

That failure can be surprising because the checkpoint file itself still exists; what changed was the code layout that Python pickling expects when reconstructing the object graph.

Another mistake is forgetting model.eval() before inference.

A third mistake is saving only model weights when the real goal was to resume training, which also requires optimizer state and other training metadata.

Summary

• Prefer torch.save(model.state_dict(), path) for model saving.
• Recreate the model class and call load_state_dict when loading.
• Use model.eval() for inference.
• Save full checkpoints when you need to resume training.
• Use map_location when moving checkpoints across devices.
• Prefer state_dict unless you have a strong reason to serialize the whole model object.