The round trip

The whole of what state_dict() puts in the dictionary is two lines of _save_to_state_dict, one for parameters and one for buffers, and both read the same way. Each entry is the live tensor, or .detach() of it, which shares the same storage and only drops the autograd history.

Chapter 3 calls this the serialization contract, and it is, once torch.save has written the bytes out. Before that write, the dictionary is a set of pointers into the model you are still training.

Sometimes the live tensors are the point, and keep_vars=True hands them over ungutted, gradients and all, for inspection rather than saving. Weight tying shows the same identity question from the other side: parameters() deduplicates by object identity while state_dict() does not, so a tied weight is one parameter under two keys, both pointing at one storage.

        for name, param in self._parameters.items():
            if param is not None:
                destination[prefix + name] = param if keep_vars else param.detach()
        for name, buf in self._buffers.items():
            if buf is not None and name not in self._non_persistent_buffers_set:
                destination[prefix + name] = buf if keep_vars else buf.detach()

Validation improves, so you keep best = model.state_dict() in a variable, training runs on, and at the end you restore from best. What you restore is the final model, because best was never holding numbers of its own.

The run below shows it in four lines: the dictionary's tensor and the parameter report the same data_ptr, and one in-place update through the parameter is visible through the dictionary immediately.

copy.deepcopy(model.state_dict()) is the fix for an in-memory snapshot, and torch.save is the fix for a durable one, since writing to disk materializes the bytes. Chapter 4 covers what else belongs in that file next to the model.

import torch
from torch import nn


class Tied(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = nn.Linear(2, 2, bias=False)
        self.b = nn.Linear(2, 2, bias=False)
        self.b.weight = self.a.weight


m = nn.Linear(2, 2)
best = m.state_dict()                                    # the "snapshot"
print(best["weight"].data_ptr() == m.weight.data_ptr())  # True
with torch.no_grad():
    m.weight.add_(1.0)                                   # one more optimizer step
print(torch.equal(best["weight"], m.weight))             # True
print(best["weight"].requires_grad,
      m.state_dict(keep_vars=True)["weight"].requires_grad)   # False True

t = Tied()
print(len(list(t.parameters())), list(t.state_dict()))   # 1 ['a.weight', 'b.weight']

Loading runs the mirror image. _load_from_state_dict finds the module's own parameter, checks the shape, and then calls param.copy_(input_param) under no_grad. The tensor object in your model is the destination and it keeps its own dtype, its own device and its own identity; only the values arrive.

That is why a float64 checkpoint loads into a float32 model without a word. The copy casts, the model stays float32, and the numbers you get are the checkpoint's rounded to single precision. Nothing in the return value mentions it.

assign=True switches the whole thing to the other semantics. Instead of copying into the existing tensor, it calls setattr with the checkpoint's tensor, so the module adopts that object whole: its dtype, its device, its storage. The measured data_ptr is the same one the checkpoint dictionary holds.

import torch
from torch import nn

ck = {"weight": torch.ones(2, 2, dtype=torch.float64),
      "bias": torch.ones(2, dtype=torch.float64)}

dst = nn.Linear(2, 2)
dst.load_state_dict(ck)
print(dst.weight.dtype, dst.weight[0, 0].item())   # torch.float32 1.0

dst2 = nn.Linear(2, 2)
dst2.load_state_dict(ck, assign=True)
print(dst2.weight.dtype,
      dst2.weight.data_ptr() == ck["weight"].data_ptr())   # torch.float64 True

Read load_state_dict and one argument in the recursive call settles a common misreading. The strict you passed is not forwarded; the literal True is, so missing and unexpected keys are collected on every load. What strict decides is whether those two lists get turned into an exception at the end.

So load_state_dict(sd, strict=False) still knows exactly which keys were missing and which were surplus. It returns them, as _IncompatibleKeys(missing_keys, unexpected_keys), and almost every call site in the wild throws that return value away. Assign it and print it and half of the load-time mysteries stop being mysteries.

One class of problem ignores strict entirely. A shape mismatch appends to error_msgs, a third list, and error_msgs raises whether or not you asked for strictness, because there is no sensible way to partially copy a differently shaped tensor. The message names the key, the checkpoint's shape and the model's, in that order.

strict=False does not make the load quieter about what it found. It makes the load stop raising about it.

import torch
from torch import nn


class Wider(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(2, 2)
        self.extra = nn.Linear(2, 2)


ck = {"lin.weight": torch.zeros(2, 2), "lin.bias": torch.zeros(2), "ghost": torch.zeros(1)}
report = Wider().load_state_dict(ck, strict=False)
print(report.missing_keys)      # ['extra.weight', 'extra.bias']
print(report.unexpected_keys)   # ['ghost']

Wider().load_state_dict({"lin.weight": torch.zeros(3, 3)}, strict=False)
# RuntimeError: Error(s) in loading state_dict for Wider:
#     size mismatch for lin.weight: copying a param with shape torch.Size([3, 3])
#     from checkpoint, the shape in current model is torch.Size([2, 2]).

A state dict carries one thing that is not a tensor. state_dict() attaches a _metadata attribute mapping each submodule prefix to a small dict, and the entry it always writes is version=self._version, a class attribute that a layer bumps when its saved state changes shape.

BatchNorm is the worked example. Its _version is 2, because version 2 added the num_batches_tracked buffer, and its own _load_from_state_dict looks at the incoming version and synthesizes that key when an older checkpoint does not have it. A version-1 checkpoint therefore loads clean under strict=True.

Stamp the same incomplete dictionary as version 2 and it fails with a missing key. The tensors did not change; the claim about what those tensors mean did. That is the mechanism available to your own modules too: bump _version, override _load_from_state_dict, migrate the old layout forward before calling super().

import torch
from torch import nn

sd = nn.BatchNorm1d(3).state_dict()
print(list(sd))
# ['weight', 'bias', 'running_mean', 'running_var', 'num_batches_tracked']
print(dict(sd._metadata))                  # {'': {'version': 2}}

old = type(sd)({k: v for k, v in sd.items() if k != "num_batches_tracked"})
old._metadata = {"": {"version": 1}}
print(nn.BatchNorm1d(3).load_state_dict(old, strict=True))
# <All keys matched successfully>

old._metadata = {"": {"version": 2}}
nn.BatchNorm1d(3).load_state_dict(old, strict=True)
# RuntimeError: Error(s) in loading state_dict for BatchNorm1d:
#     Missing key(s) in state_dict: "num_batches_tracked".

The assign argument reaches _load_from_state_dict through that same metadata, and the line that puts it there writes into the dictionary the caller owns rather than into a copy. local_metadata is the per-prefix dict stored inside state_dict._metadata, so local_metadata['assign_to_params_buffers'] = assign is a permanent edit to the checkpoint object.

Load once with assign=True and every later load from that same dictionary assigns, whether or not you ask. Two models loaded from it afterwards do not get copies; they get the same tensors, sharing one storage, so a training step on one moves the other.

The line is unchanged upstream at v2.8.0, at module.py:2571-2575, so this is a property of the API rather than a fixed bug in an old release. The safe habit is one dictionary per load when assign is in play, or a fresh state_dict() call, which builds a fresh _metadata every time.

import torch
from torch import nn

ck = nn.Linear(4, 4).state_dict()
print(dict(ck._metadata))                  # {'': {'version': 1}}

nn.Linear(4, 4, device="meta").load_state_dict(ck, assign=True)
print(dict(ck._metadata))
# {'': {'version': 1, 'assign_to_params_buffers': True}}

a, b = nn.Linear(4, 4), nn.Linear(4, 4)
a.load_state_dict(ck)                      # no assign asked for
b.load_state_dict(ck)
print(a.weight.data_ptr() == b.weight.data_ptr())   # True
with torch.no_grad():
    a.weight.add_(1.0)
print(torch.equal(a.weight, b.weight))              # True