Saving and Loading Models in PyTorch: state_dict, Checkpoints and Resuming Training

A trained model is saved for one of two reasons. The first is to use it later for predictions, which needs only the model's weights. The second is to continue training later, after a crash, a disconnected notebook or a planned stop, which needs the weights and also the optimizer's internal state. This article covers both: what a state_dict is, how to save and load one, how to keep the best model during training, and what a full checkpoint must contain for training to resume exactly where it stopped.

What a state_dict is

model.state_dict() returns an ordered dictionary. The keys are the names of the model's tensors and the values are the tensors themselves.

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 3))

for k, v in model.state_dict().items():
    print(f"{k:12s} {tuple(v.shape)}")

Output:

0.weight     (32, 10)
0.bias       (32,)
2.weight     (3, 32)
2.bias       (3,)

The name is built from the position of the layer in the model and the name of the tensor inside the layer. 0.weight is the weight of the layer at index 0. The ReLU at index 1 has no tensors, so it has no entries. In a model written as a class, the names come from the attribute names, so a layer stored as self.fc1 gives the keys fc1.weight and fc1.bias.

A state_dict holds more than the learnable parameters. It also holds the module's buffers, which are tensors that belong to the model but are not trained by the optimizer. BatchNorm keeps its running statistics in buffers:

bn_model = nn.Sequential(nn.Linear(10, 32), nn.BatchNorm1d(32), nn.ReLU(), nn.Linear(32, 3))

for k, v in bn_model.state_dict().items():
    print(f"{k:28s} {tuple(v.shape)}")

0.weight                     (32, 10)
0.bias                       (32,)
1.weight                     (32,)
1.bias                       (32,)
1.running_mean               (32,)
1.running_var                (32,)
1.num_batches_tracked        ()
3.weight                     (3, 32)
3.bias                       (3,)

This model has 6 tensors in parameters() and 9 entries in its state_dict. The three extra entries are the BatchNorm buffers. They are needed to reproduce the model's evaluation-mode output, so a saved model without them would not predict the same way.

Saving and loading the weights

Saving is one call:

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

Loading has two steps. First create a model with the same architecture, then copy the saved tensors into it:

model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 3))
model.load_state_dict(torch.load("model.pt", weights_only=True))
model.eval()

The file holds only tensors and their names. It does not hold the architecture, so the code that defines the model is needed again at load time. model.eval() after loading puts Dropout and BatchNorm layers into evaluation mode before the model is used for predictions.

weights_only=True tells torch.load to accept only tensors and simple Python types (dicts, lists, tuples, numbers, strings) and to refuse anything else. Since PyTorch 2.6 this is the default, so torch.load("model.pt") behaves the same way. The reason for the restriction is explained below, in the section on saving the whole model.

A checkpoint saved on a GPU loads on a machine without one when the tensors are moved at load time:

state = torch.load("model.pt", map_location="cpu", weights_only=True)

state_dict() returns references, not copies

The tensors in a state_dict are the model's own tensors, not copies of them. When the model is trained further, the values in a previously returned state_dict change with it:

m = nn.Linear(4, 1)
sd = m.state_dict()
print(sd["weight"].data_ptr() == m.weight.data_ptr())

before = sd["weight"].clone()
with torch.no_grad():
    m.weight.add_(1.0)
print((sd["weight"] - before).abs().max().item())

True
1.0

The first line shows that the tensor in sd and the model's weight use the same memory. The second shows that sd changed when the weight was updated, although sd itself was never touched.

torch.save writes the values to disk at the moment it is called, so the saved file does not change afterwards. The problem appears only when a state_dict is kept in a Python variable as a "saved" copy. To keep a copy in memory, use copy.deepcopy:

import copy
best_state = copy.deepcopy(model.state_dict())

Saving the best model during training

The last epoch is often not the best one. If the model starts to overfit, the validation loss reaches a minimum and then rises. The standard pattern keeps the lowest validation loss seen so far and saves the weights each time it improves:

best_val = float("inf")

for epoch in range(num_epochs):
    # ... training phase ...
    # ... validation phase, gives avg_val ...
    if avg_val < best_val:
        best_val = avg_val
        torch.save(model.state_dict(), "best.pt")

The example below is a regression network trained with Adam on 60 samples and validated on 240, which is small enough to overfit:

epoch   0  train 9.6926  val 6.1876
epoch  10  train 7.0336  val 4.9666
epoch  20  train 3.7911  val 3.1780
epoch  30  train 1.0966  val 1.4995
epoch  40  train 0.5580  val 1.0054
epoch  60  train 0.2154  val 1.0248
epoch 100  train 0.0506  val 1.0713
epoch 150  train 0.0079  val 1.1214
epoch 199  train 0.0010  val 1.1498

The lowest validation loss, 0.9740, was reached at epoch 46. After that the training loss kept falling while the validation loss rose to 1.1498 at epoch 199. In the same run, the best weights were also kept in memory in two ways, once with deepcopy and once by assigning best_ref = model.state_dict(). Each was then loaded into a fresh model and evaluated:

reload best.pt              val 0.9740
deepcopy kept in memory     val 0.9740
plain state_dict() kept     val 1.1498

The file and the deep copy hold the epoch-46 weights. The plain state_dict() reference followed the model to the end of training and holds the epoch-199 weights, which is why it gives the final validation loss and not the best one.

Errors when loading

load_state_dict matches entries by name and checks every shape. By default (strict=True) any difference raises a RuntimeError that lists all the problems. The examples load the state_dict of Sequential(Linear(10, 32), ReLU(), Linear(32, 3)) into different models.

Different layer size. A hidden layer of 64 instead of 32:

RuntimeError: Error(s) in loading state_dict for Sequential:
    size mismatch for 0.weight: copying a param with shape torch.Size([32, 10]) from checkpoint, the shape in current model is torch.Size([64, 10]).
    size mismatch for 0.bias: copying a param with shape torch.Size([32]) from checkpoint, the shape in current model is torch.Size([64]).
    size mismatch for 2.weight: copying a param with shape torch.Size([3, 32]) from checkpoint, the shape in current model is torch.Size([3, 64]).

An extra layer in the model. Two more layers added at the end:

RuntimeError: Error(s) in loading state_dict for Sequential:
    Missing key(s) in state_dict: "4.weight", "4.bias".

"Missing" means the model has a tensor that the file does not provide. "Unexpected" means the file has a tensor that the model does not have.

Same layers, different names. The same two Linear layers in a class that stores them as self.fc1 and self.fc2:

RuntimeError: Error(s) in loading state_dict for Net:
    Missing key(s) in state_dict: "fc1.weight", "fc1.bias", "fc2.weight", "fc2.bias".
    Unexpected key(s) in state_dict: "0.weight", "0.bias", "2.weight", "2.bias".

The shapes match, but the names do not, and names are what load_state_dict uses. A model rewritten from nn.Sequential to a class, or with renamed attributes, cannot load an old file directly. The keys have to be renamed in the dictionary first.

load_state_dict(sd, strict=False) does not raise. It loads the entries whose names match and returns the ones that did not:

_IncompatibleKeys(missing_keys=['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias'], unexpected_keys=['0.weight', '0.bias', '2.weight', '2.bias'])

In this case nothing matched, so nothing was loaded and the model keeps its random initial weights. strict=False is meant for loading part of a model on purpose, for example a pretrained body with a new final layer. When it is used, the returned lists should be checked.

Why the state_dict and not the whole model

torch.save(model, "full.pt") also works. It saves the model object with Python's pickle, which stores a reference to the model's class (its module and name) and the object's attributes. Loading it has two requirements. The class definition must be importable under the same module path, so renaming the file or moving the class breaks loading. And unpickling can run arbitrary code, so a file from an unknown source can execute anything when loaded.

The default weights_only=True therefore refuses it:

torch.save(net, "full.pt")
torch.load("full.pt")

UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint. ...

torch.load("full.pt", weights_only=False) loads it, and should only be used for files you created yourself. A state_dict file avoids both problems: it contains only names and tensors, it loads with weights_only=True, and it can be loaded into any model whose names and shapes match.

The optimizer's state_dict

Optimizers have a state_dict too. For Adam:

opt = torch.optim.Adam(net.parameters())
# ... one training step ...
od = opt.state_dict()
print(list(od.keys()))
print(list(od["state"][0].keys()))
print(od["param_groups"][0]["params"])

['state', 'param_groups']
['step', 'exp_avg', 'exp_avg_sq']
[0, 1, 2, 3]

param_groups holds the settings: learning rate, betas, weight decay and so on. Parameters are identified by position (0, 1, 2, 3), not by name, so the optimizer must be created over the same parameters in the same order before its state is loaded.

state holds, for each parameter, the step count and Adam's two running averages: exp_avg (the average of the gradients) and exp_avg_sq (the average of the squared gradients). Each running average has the same shape as its parameter, so Adam's state is twice the size of the model. For the small Net above, the model has 451 numbers and Adam's buffers hold 902. For a 1,409-parameter model, the file with weights only is 7,649 bytes, and the file with weights and Adam's state is 22,259 bytes.

A full checkpoint

To resume training, save everything the training loop depends on in one dictionary:

checkpoint = {
    "epoch": epoch,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "best_val": best_val,
    "history": history,
}
torch.save(checkpoint, "checkpoint.pt")

And to resume:

ckpt = torch.load("checkpoint.pt", weights_only=True)

model = MyModel()
model.load_state_dict(ckpt["model"])
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
optimizer.load_state_dict(ckpt["optimizer"])

start_epoch = ckpt["epoch"] + 1
best_val    = ckpt["best_val"]
history     = ckpt["history"]

for epoch in range(start_epoch, num_epochs):
    ...

This dictionary loads with weights_only=True, because everything in it is a tensor, a dict, a list or a Python number. It loads with no error:

checkpoint dict with weights_only=True: loaded, keys = ['epoch', 'model', 'optimizer', 'best_val', 'history']

The restriction applies to the types inside the dictionary. If best_val is a NumPy number instead of a Python float, for example because it came from np.mean, loading fails:

UnpicklingError: Weights only load failed. ...
Unsupported global: GLOBAL numpy._core.multiarray.scalar was not an allowed global by default. ...

Converting values with float(...) or .item() before saving keeps the checkpoint loadable with weights_only=True.

What happens without the optimizer state

If only the model is restored and a new optimizer is created, Adam starts again with step = 0 and both running averages at zero. On its first step, Adam's update for each weight is approximately lr * g / |g|, which is lr in size for every weight whatever its gradient. A trained model normally takes much smaller steps, because its running averages describe gradients that are small and noisy.

Measured on the first layer of a model restored after 10 epochs of training, with lr = 0.01:

restored Adam: first step, mean |change| of layer-0 weights 0.001007   max 0.006382
fresh Adam:    first step, mean |change| of layer-0 weights 0.009999   max 0.010000

The fresh optimizer moves every weight by 0.01, about ten times the average step of the restored one. On the overfitted regression model from earlier, which had reached a training loss of 0.0009, five steps after resuming gave:

restored optimizer   train loss: 0.0009, 0.0009, 0.0009, 0.0008, 0.0008, 0.0007
fresh optimizer      train loss: 0.0009, 0.0016, 0.0561, 0.0009, 0.0293, 0.0349

With the restored state the loss keeps decreasing. With a fresh optimizer it jumps to 0.0561 after two steps, about 60 times higher. Training usually recovers, but the first epochs after resuming are spent undoing these steps.

Resuming exactly: random number state

Model and optimizer state are enough to continue training. For the resumed run to be identical to an uninterrupted run, the random number generators also have to be restored, because they decide the shuffle order of the DataLoader and which units Dropout removes.

g = torch.Generator()
loader = DataLoader(dataset, batch_size=64, shuffle=True, generator=g)

# when saving
checkpoint["torch_rng"]  = torch.get_rng_state()
checkpoint["loader_rng"] = g.get_state()

# when resuming
torch.set_rng_state(ckpt["torch_rng"])
g.set_state(ckpt["loader_rng"])

A classifier with Dropout was trained for 20 epochs without stopping, and also stopped at epoch 10 and resumed three ways:

epoch  no stop   opt + rng   opt, no rng   no opt (rng restored)
 10    0.02206   0.02206     0.02573       0.02781
 11    0.02069   0.02069     0.02299       0.02101
 12    0.02075   0.02075     0.02442       0.02543
 13    0.02263   0.02263     0.02317       0.02703
 14    0.02326   0.02326     0.01231       0.02531
 15    0.02584   0.02584     0.01787       0.02472
 16    0.01977   0.01977     0.02367       0.01787
 17    0.01713   0.01713     0.01703       0.01609
 18    0.02019   0.02019     0.01730       0.01817
 19    0.02110   0.02110     0.01510       0.02156

With the optimizer and the random state restored, every loss matches the uninterrupted run, and the final weights are bit-for-bit equal (torch.equal is True for every tensor). Without the random state, the run continues normally but takes a different path, and its final weights differ from the uninterrupted run by up to 0.40. On a GPU, torch.cuda.get_rng_state_all() and torch.cuda.set_rng_state_all() save and restore the GPU generators in the same way.

Exact reproduction is not always needed. It is useful when comparing two runs, or when checking that resuming did not change the result.

Logging the numbers to disk

A history dictionary in memory is lost when the process stops. Appending one row per epoch to a CSV file keeps the numbers on disk:

import csv

with open("log.csv", "w", newline="") as f:
    csv.writer(f).writerow(["epoch", "train_loss", "val_loss"])

for epoch in range(num_epochs):
    # ... training + validation ...
    with open("log.csv", "a", newline="") as f:
        csv.writer(f).writerow([epoch, f"{avg_train:.6f}", f"{avg_val:.6f}"])

The first lines of the file from the regression run:

epoch,train_loss,val_loss
0,9.692556,6.187602
1,9.419966,6.079378
2,9.157091,5.970488

The file is opened and closed every epoch, so each row is written to disk as soon as the epoch ends. Tools such as TensorBoard, Weights & Biases and MLflow do the same job with plots and run comparison built in.

How often to save

Strategy	What is saved	Use
Best model only	best.pt, overwritten when validation loss improves	the model to use after training
Latest checkpoint	checkpoint.pt, overwritten every N epochs	resuming after a crash
Numbered checkpoints	ckpt_010.pt, ckpt_020.pt, ...	going back to an earlier point, at the cost of disk space

Most projects use the first two together: best.pt with weights only for the final model, and one full checkpoint that is overwritten regularly for resuming.

Summary

• model.state_dict() maps tensor names to tensors. It contains the parameters and the buffers, such as BatchNorm's running statistics.
• Save with torch.save(model.state_dict(), path). Load by creating the same architecture and calling model.load_state_dict(torch.load(path, weights_only=True)), then model.eval() for predictions.
• state_dict() returns references to the model's tensors. Use copy.deepcopy to keep a copy in memory.
• load_state_dict matches by name and shape. Size mismatches, missing keys and unexpected keys raise an error unless strict=False, which returns the keys that did not match.
• Save the state_dict, not the model object. Pickled models depend on the class location and need weights_only=False, which should only be used for trusted files.
• To resume training, save the model, the optimizer, the epoch, the best validation loss and the history. Without the optimizer state, Adam's first steps are as large as the learning rate for every weight.
• For a resumed run to match an uninterrupted one exactly, also save and restore the random number state.

---

This is one chapter's worth of an idea from my book, PyTorch From Ground Up, which builds everything from tensors upward so nothing stays vague. If it helped: 8 chapters are free, no email required, there's a free one-page tensor cheat-sheet here, every example runs in the companion notebooks on GitHub, and the full book is on Leanpub or in paperback and Kindle on Amazon.

---

More in this series

How Training Actually Works:

• PyTorch Autograd Explained: What .backward() Actually Does
• Backpropagation by Hand: Two Layers, a Pen, and Then Autograd Agrees
• What optimizer.step() Actually Does: SGD, Momentum and Adam by Hand
• What Your Loss Function Actually Tells the Model
• nn.Module Explained: The Same Model Built with Raw Tensors and with nn.Module
• model.train() vs model.eval() vs torch.no_grad() in PyTorch

Shape mechanics:

• Reshape vs View in PyTorch
• PyTorch Broadcasting Explained
• What Does unsqueeze Do in PyTorch?
• What Does keepdim Do in PyTorch?
• PyTorch permute vs transpose

Next in this series: moving a model and its data to the GPU.