PyTorch From First Principles
Chapter 11 of 16

nn.Module: What Does PyTorch Think Belongs to Your Model?

By AI, PyTorch

Explain this chapter with AI

Copy this prompt into ChatGPT, Claude, Gemini, a local model, or another AI.

Apply this chapter with AI

Copy this prompt into ChatGPT, Claude, Gemini, a local model, or another AI.

Here is a classifier for the XOR-shaped dataset from Chapter 4, written the way most people write PyTorch. It trains. Validation accuracy climbs from chance to 95.7%. It saves a checkpoint, the checkpoint loads without a warning, and then the reloaded model performs at chance.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

n = 2000
X = torch.randn(n, 2)
y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).long()

perm = torch.randperm(n)
X_train, y_train = X[perm[:1600]], y[perm[:1600]]
X_val, y_val = X[perm[1600:]], y[perm[1600:]]

class Encoder(nn.Module):
    def __init__(self, width=16, depth=3):
        super().__init__()
        self.blocks = [nn.Linear(width, width) for _ in range(depth)]

    def forward(self, x):
        for block in self.blocks:
            x = torch.relu(block(x))
        return x

class Classifier(nn.Module):
    def __init__(self, width=16, depth=3):
        super().__init__()
        self.stem = nn.Linear(2, width)
        self.encoder = Encoder(width, depth)
        self.head = nn.Linear(width, 2)

    def forward(self, x):
        x = torch.relu(self.stem(x))
        x = self.encoder(x)
        return self.head(x)

model = Classifier()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

for epoch in range(200):
    logits = model(X_train)
    loss = F.cross_entropy(logits, y_train)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 50 == 0 or epoch == 199:
        with torch.no_grad():
            val_acc = (model(X_val).argmax(1) == y_val).float().mean()
        print(f"epoch={epoch:3d} loss={loss.item():.4f} "
              f"val_acc={val_acc.item():.3f}")

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

reloaded = Classifier()
print(reloaded.load_state_dict(torch.load("classifier.pt", weights_only=True)))

with torch.no_grad():
    acc = (reloaded(X_val).argmax(1) == y_val).float().mean()
print(f"reloaded val_acc={acc.item():.3f}")

epoch=  0 loss=0.6964 val_acc=0.512
epoch= 50 loss=0.5930 val_acc=0.810
epoch=100 loss=0.3856 val_acc=0.873
epoch=150 loss=0.2550 val_acc=0.925
epoch=199 loss=0.1908 val_acc=0.957
<All keys matched successfully>
reloaded val_acc=0.512

Read the last two lines together. PyTorch reports that every key in the checkpoint matched every key the model expected, and the model that came back is worthless. Nothing raised. Nothing warned. load_state_dict was as satisfied as it is capable of being.

The tools from the previous three chapters do not find this. Every shape is right. Every dtype and device is right. The gradient reaches every weight in the network, and we could prove it. The loss falls monotonically, which is the outcome Chapter 1 taught us to want, and validation accuracy rises with it, which is the check Chapter 1 taught us to add. The six diagnostic questions from Chapter 4 — representation, computation, objective, differentiation, update, evaluation — are all things we could ask here. Most of the evidence looks healthy. The computation works, the objective is valid, gradients flow, and evaluation improves.

The failure appears only when we ask a different question: which of the tensors participating in that computation does PyTorch actually manage as state belonging to this model?

One number gives it away:

print("parameter tensors:", len(list(model.parameters())))
print("parameter count:  ", sum(p.numel() for p in model.parameters()))

parameter tensors: 4
parameter count:   82

Eighty-two. Chapter 4’s hand-built two-layer network had exactly eighty-two parameters, and this model has five nn.Linear layers. Three of them — 816 parameters — are not in that count. They are not in the checkpoint either, which is why the checkpoint loaded so cleanly: the model that loaded it did not think it was missing anything.

Those three layers exist. They are constructed, they are called on every forward pass, they transform the data, and autograd differentiates through them. What has gone wrong is not in the mathematics of the model. It is that PyTorch and the programmer disagree about what the model is.

Chapter 4 ended by asking three questions about the framework we were about to meet: why model.parameters() finds a weight buried four levels deep, why model.to("cuda") moves tensors you never named, and why a module stored in a plain Python list can vanish from state_dict(). They have one answer between them, and we have just watched the third one cost us a trained model.

That answer is recursive composition — modules containing modules, and framework operations walking the result — but stated that way it is a slogan, and slogans do not survive a debugging session. What makes it useful is knowing exactly what the walk can see:

Which objects does PyTorch consider part of this model, which piece of machinery consults that answer, and how do I check it instead of inferring it from the source code?

By the end you should be able to read an unfamiliar model class and predict what named_parameters(), state_dict(), .to(device), eval() and an optimizer will each find in it — and when one of them finds something you did not expect, know which structure to inspect first.

Four systems, four different answers

Before repairing anything, get the evidence. The first question is what the framework thinks it is holding.

for name, module in model.named_modules():
    print(f"{name!r:12s} {type(module).__name__}")

''           Classifier
'stem'       Linear
'encoder'    Encoder
'head'       Linear

Four modules, in a network that contains seven. The three Linear layers inside Encoder are absent. So are their parameters:

for name, p in model.named_parameters():
    print(f"{name:12s} {tuple(p.shape)}")

stem.weight  (16, 2)
stem.bias    (16,)
head.weight  (2, 16)
head.bias    (2,)

Now the second question, which is a different one: does autograd reach those layers? Take a snapshot, train briefly, and check both the gradients and the values.

before = [b.weight.detach().clone() for b in model.encoder.blocks]

optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
for _ in range(50):
    loss = F.cross_entropy(model(X_train), y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

for i, (old, block) in enumerate(zip(before, model.encoder.blocks)):
    moved = (block.weight - old).abs().max().item()
    print(f"encoder.blocks[{i}].weight  max_change={moved:.6f}  "
          f"grad_is_none={block.weight.grad is None}")

encoder.blocks[0].weight  max_change=0.000000  grad_is_none=False
encoder.blocks[1].weight  max_change=0.000000  grad_is_none=False
encoder.blocks[2].weight  max_change=0.000000  grad_is_none=False

Gradients are present. Values never move. Not approximately, not slowly: the maximum absolute change across fifty optimizer steps is exactly zero.

This is the pattern the chapter is built around, and it is worth stating plainly because it contradicts the reflex most people bring to PyTorch:

A gradient tells you that autograd reached the tensor. It tells you nothing about whether anything is going to update it.

Three separate systems have now given three different answers about the same three layers. Ordinary Python reaches them, because self.blocks is a list and the for loop in forward iterates it. Autograd reaches them, because they participated in the computation that produced the loss. The module registry does not have them, so model.parameters() does not return them, so the optimizer was never given them, so optimizer.step() has nothing to do with them.

There is a fourth symptom hiding in that loop, and it is worth seeing because it is a genuine resource bug rather than only a correctness one:

for step in range(5):
    loss = F.cross_entropy(model(X_train), y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print(f"step={step} "
          f"stem.weight.grad_norm={model.stem.weight.grad.norm().item():.4f} "
          f"blocks[0].weight.grad_norm="
          f"{model.encoder.blocks[0].weight.grad.norm().item():.4f}")

step=0 stem.weight.grad_norm=0.0048 blocks[0].weight.grad_norm=0.0145
step=1 stem.weight.grad_norm=0.0050 blocks[0].weight.grad_norm=0.0287
step=2 stem.weight.grad_norm=0.0054 blocks[0].weight.grad_norm=0.0432
step=3 stem.weight.grad_norm=0.0059 blocks[0].weight.grad_norm=0.0580
step=4 stem.weight.grad_norm=0.0062 blocks[0].weight.grad_norm=0.0734

The registered parameter’s gradient is cleared by optimizer.zero_grad() on every iteration. The unregistered parameter’s gradient is not, so each backward pass adds another contribution to the value already stored in .grad. Chapter 1 established that PyTorch accumulates gradients into .grad and that something has to clear them; optimizer.zero_grad() is that something, and it can only clear the parameters it was given. Everything else accumulates for the entire run, holding a gradient buffer the size of the weight and never using it.

None of this is PyTorch behaving inconsistently. Each system answered the question it was asked. The problem is that we assumed those questions have the same answer.

What self.something = ... actually does

The disagreement starts at assignment. nn.Module overrides __setattr__, so assigning an attribute on a module is not the ordinary Python operation it looks like. Depending on the type of the value, it goes to one of four places, and this is directly observable:

class Probe(nn.Module):
    def __init__(self):
        super().__init__()
        self.child = nn.Linear(2, 2)
        self.weight = nn.Parameter(torch.zeros(3))
        self.plain = torch.zeros(3)
        self.register_buffer("running", torch.zeros(3))
        self.number = 7

p = Probe()

print("ordinary __dict__:", [k for k in p.__dict__ if not k.startswith("_")])
print("_parameters:      ", list(p._parameters))
print("_buffers:         ", list(p._buffers))
print("_modules:         ", list(p._modules))

ordinary __dict__: ['training', 'plain', 'number']
_parameters:       ['weight']
_buffers:          ['running']
_modules:          ['child']

Five assignments, four destinations. An nn.Module value goes into _modules. An nn.Parameter value goes into _parameters. A buffer, which has to be declared explicitly because a plain tensor has no way to signal the intent, goes into _buffers via register_buffer. Anything else — including an ordinary tensor, and including a Python list that happens to contain modules — lands in the instance __dict__ like any normal Python attribute.

Attribute access is arranged so that you cannot tell the difference. p.weight returns the object in _parameters, p.child returns the object in _modules, and they are the same objects, not copies:

print(p.weight is p._parameters["weight"])
print(p.child is p._modules["child"])
print(p.running is p._buffers["running"])

True
True
True

So self.blocks = [nn.Linear(16, 16), ...] is not a failed registration. It is a successful ordinary attribute assignment. The list is stored, the modules inside it are alive and callable, and PyTorch has no reason to look inside a list it was handed as opaque Python data. Those three dictionaries are the core registries that describe the module’s parameters, buffers and child modules. Other bookkeeping exists around them, but these are the structures that matter for the ownership questions in this chapter.

Two consequences of this mechanism are worth knowing before they surprise you.

The dictionaries do not exist until Module.__init__ runs. This is why the super().__init__() line is not optional boilerplate:

class NoSuper(nn.Module):
    def __init__(self):
        self.w = nn.Parameter(torch.zeros(3))

NoSuper()

AttributeError: cannot assign parameters before Module.__init__() call

Registration is sticky by name. Once weight is in _parameters, assigning a plain tensor to that name is refused rather than silently demoting it:

layer = nn.Linear(2, 2)
layer.weight = torch.zeros(2, 2)

TypeError: cannot assign 'torch.FloatTensor' as parameter 'weight'
(torch.nn.Parameter or None expected)

That is a useful guardrail, and it is worth noticing how narrow it is. It fires only for a name that is already registered. A brand-new name gets no such protection, which is exactly the case in our broken model.

Names are paths through the tree

The three dictionaries hold direct children only. Classifier._modules contains encoder; it knows nothing about what is inside encoder. The dotted names you see in named_parameters() and state_dict() come from a recursive walk, and the walk is short enough to write out in full.

def walk_parameters(module, prefix=""):
    for name, param in module._parameters.items():
        if param is not None:
            yield prefix + name, param

    for name, child in module._modules.items():
        if child is not None:
            yield from walk_parameters(child, prefix + name + ".")

Run it against the repaired model we are about to build, and compare object identity with what PyTorch produces:

mine = [(n, id(p)) for n, p in walk_parameters(model)]
theirs = [(n, id(p)) for n, p in model.named_parameters()]

for name, _ in mine:
    print(name)

print()
print("identical to named_parameters():", mine == theirs)

stem.weight
stem.bias
encoder.blocks.0.weight
encoder.blocks.0.bias
encoder.blocks.1.weight
encoder.blocks.1.bias
encoder.blocks.2.weight
encoder.blocks.2.bias
head.weight
head.bias

identical to named_parameters(): True

For this model, the miniature walk produced the same names in the same observed order and pointed at the same parameter objects as named_parameters().

Do not treat the ordering as a contract — PyTorch’s real traversal handles additional cases such as duplicate objects and other bookkeeping. The important result is structural: registered names are paths through a recursively composed module hierarchy.

That single sentence explains a large fraction of the surprising behavior in this chapter. encoder.blocks.0.weight is not a label PyTorch invented for the layer. It is a route: attribute encoder on the root, attribute blocks on that, key 0 on that, parameter weight on that. Rename self.encoder to self.backbone and every one of those names changes, because the route changed. Move a layer from one container to another and the names change. This is why checkpoints are sensitive to refactoring in a way that surprises people, and we will come back to it.

It also explains why the parent needs no special code. Classifier never mentions Linear.weight. It registers encoder, Encoder registers blocks, ModuleList registers its entries, and Linear registers its own weight and bias. Each level knows only its direct children, and the recursion assembles the rest. A transformer with a hundred blocks is not a harder case than this one; it is the same walk, further.

The four structures

We can now name the four structures precisely, because we have seen three of them disagree.

The Python object graph. Whatever ordinary attribute access can reach: lists, dicts, closures, module-level globals, objects held by other objects. This is what forward() uses. It is the largest of the four and PyTorch does not traverse it.

The registered module tree. The _modules, _parameters and _buffers dictionaries, walked recursively. This is a declaration of ownership, built entirely by assignment, and it is what the framework traverses.

The autograd graph. Built fresh during each forward pass from the operations that actually executed, as Chapter 3 described. It records how this particular loss depended on which tensors. It has no concept of a model at all.

The optimizer’s parameter groups. A list of tensor objects handed to a specific optimizer at construction time, stored by reference. It is a snapshot of a decision, not a live view of anything.

The same tensor can participate in several of these structures at once, and the memberships do not have to agree. The three encoder layers in our broken model are in the first and third and neither of the others. Once you see the four as separate, the framework stops looking like a collection of rules and starts looking like a set of traversals, and every API becomes a question about which structure it walks:

named_modules()      → registered module tree
named_parameters()   → registered parameters, found by walking that tree
named_buffers()      → registered buffers, found by walking that tree
state_dict()         → registered parameters and persistent buffers, keyed by path
model.to(...)        → registered parameters and buffers
model.train/eval()   → the `training` flag on every registered module
loss.backward()      → the autograd graph recorded by this forward pass
optimizer.step()     → the tensors in this optimizer's parameter groups

When something is missing, the productive question is not “why did PyTorch ignore my tensor” but:

Which of these four structures does the failing operation traverse, and is my object in it?

That is this chapter’s addition to the diagnostic method. Chapter 2 said find the first wrong tensor. Chapter 3 said find where the gradient path first disappears. Chapter 4 said determine which of the six claims is failing. This chapter says: when something appears to vanish, compare the structure you intended against the structure PyTorch actually registered, and do it by inspection rather than by re-reading the source.

A tensor with a gradient that is not a parameter

The plain-list bug hides three whole layers, which makes it dramatic but also easy to dismiss as carelessness. The same mechanism produces a much smaller failure that is harder to see. Here is a linear layer written by hand:

class BrokenLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = torch.randn(out_features, in_features) * 0.5
        self.weight.requires_grad_(True)
        self.bias = torch.zeros(out_features, requires_grad=True)

    def forward(self, x):
        return x @ self.weight.T + self.bias

Before running anything, predict the answers. This is worth doing on paper; the value of the experiment is in finding out which of your predictions were coupled together when they should not have been.

weight.requires_grad          ?
weight.is_leaf                ?
weight.grad after backward    ?
model.named_parameters()      ?
model.state_dict()            ?
does model.to(...) convert it ?
can an optimizer update it    ?

Now run it:

model = BrokenLinear(4, 2)

x = torch.randn(8, 4)
y = torch.randint(0, 2, (8,))
F.cross_entropy(model(x), y).backward()

w = model.weight
print("requires_grad   :", w.requires_grad)
print("is_leaf         :", w.is_leaf)
print("grad is None    :", w.grad is None)
print("grad norm       :", round(w.grad.norm().item(), 4))
print("named_parameters:", list(model.named_parameters()))
print("state_dict keys :", list(model.state_dict().keys()))

model.to(torch.float64)
print("dtype after .to :", model.weight.dtype)

requires_grad   : True
is_leaf         : True
grad is None    : False
grad norm       : 0.6618
named_parameters: []
state_dict keys : []
dtype after .to : torch.float32

Everything Chapter 3 cares about is in perfect order. weight is a leaf, it requires gradients, backward() populated .grad with a finite value, and the arithmetic in forward is correct. Everything this chapter cares about is empty. The model has no parameters, saves nothing, and does not convert.

The last of the seven predictions produces the one loud failure in this example:

torch.optim.SGD(model.parameters(), lr=0.1)

ValueError: optimizer got an empty parameter list

Worth pausing on. PyTorch does check, and it does complain — but only about a model with zero registered parameters. Our opening Classifier had four, so the check passed and the run proceeded in silence. Loud failures are the fortunate case. The dangerous version of this bug is always the partial one.

nn.Parameter is a declaration, not a flavor of tensor

The repair is one word in two places:

class ManualLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features) * 0.5)
        self.bias = nn.Parameter(torch.zeros(out_features))

    def forward(self, x):
        return x @ self.weight.T + self.bias

The same probes, plus two that were not worth running before:

type            : Parameter
requires_grad   : True
is_leaf         : True
grad is None    : False
named_parameters: ['weight', 'bias']
state_dict keys : ['weight', 'bias']
dtype after .to : torch.float64
optimizer holds : 2 tensors

The first four lines are identical to the broken version. That is the point. nn.Parameter did not make the tensor differentiable — it already was. It did not make it a leaf — it already was. What changed is which dictionary the assignment landed in, and therefore what the recursive walk can find.

This is why the common gloss is worth refusing:

nn.Parameter is a tensor with requires_grad=True.

It is false in both directions. Chapter 3 gave us plenty of tensors with requires_grad=True that are not parameters of anything. And a parameter does not have to require gradients:

class Frozen(nn.Module):
    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.ones(3), requires_grad=False)

f = Frozen()
print("registered   :", [n for n, _ in f.named_parameters()])
print("requires_grad:", f.w.requires_grad)
print("in state_dict:", list(f.state_dict().keys()))

registered   : ['w']
requires_grad: False
in state_dict: ['w']

A registered parameter that autograd will not track. Fully in the model, fully in the checkpoint, follows the model across devices, and receives no gradient. Nothing about that is contradictory once you accept that registration and gradient tracking are answers to different questions:

assign nn.Parameter to a Module attribute
        → "register this object as a parameter of this module"

requires_grad
        → "record differentiable operations involving this tensor when grad mode permits it"

There is a sharper version of the first line. Being an nn.Parameter is not a property that makes a tensor belong to a model; registration is a relationship between an object and a particular module. In our broken Classifier, model.encoder.blocks[0].weight is a genuine nn.Parameter — nn.Linear created it that way — and it is registered on the Linear that owns it. It is simply not reachable from model by the walk, because the Linear itself was never registered on Encoder. Asking “is this an nn.Parameter?” does not answer “is this in model.parameters()?” You have to ask the model.

Repairing the model, and the evidence that it worked

Back to the opening. The three encoder layers need to be registered children of Encoder, and the container that holds them needs to be one PyTorch traverses. One line changes:

class Encoder(nn.Module):
    def __init__(self, width=16, depth=3):
        super().__init__()
        self.blocks = nn.ModuleList(
            [nn.Linear(width, width) for _ in range(depth)]
        )

    def forward(self, x):
        for block in self.blocks:
            x = torch.relu(block(x))
        return x

forward is untouched. The arithmetic is untouched. The initialization draws the same numbers from the same seed. Now rerun the identical script and compare it against the identical measurements.

''                   Classifier
'stem'               Linear
'encoder'            Encoder
'encoder.blocks'     ModuleList
'encoder.blocks.0'   Linear
'encoder.blocks.1'   Linear
'encoder.blocks.2'   Linear
'head'               Linear

parameter tensors: 10
parameter count:   898

epoch=  0 loss=0.6964 val_acc=0.512
epoch= 50 loss=0.0185 val_acc=0.995
epoch=100 loss=0.0081 val_acc=0.995
epoch=150 loss=0.0059 val_acc=0.995
epoch=199 loss=0.0033 val_acc=0.995
<All keys matched successfully>
reloaded val_acc=0.995

Put the two runs side by side:

measurement plain list nn.ModuleList
registered modules 4 8
len(list(model.parameters())) 4 10
parameter count 82 898
state_dict() keys 4 10
REMOVE THIS ROW
epoch 0 loss 0.6964 0.6964
final validation accuracy 0.957 0.995
accuracy after save and reload 0.512 0.995

The epoch-0 row is the one to dwell on. Identical loss, to four decimal places, because the two programs perform exactly the same computation on exactly the same numbers. The entire difference between the two columns is what the framework was told about that computation.

The final-accuracy row is the trap. The broken model reached 95.7%, which is a respectable number that no monitoring dashboard would flag. It got there because the stem and head were trainable and the frozen random middle layers acted as a fixed nonlinear feature map — a legitimate model, just not the one anybody wrote down. The failure only became visible when the model had to survive leaving the process.

And the last row is the reason this class of bug deserves a chapter. The checkpoint was not corrupt. It faithfully recorded every piece of state PyTorch believed the model had.

Reachable in Python is not the same as registered in PyTorch. Used in the computation is not the same as managed as model state.

That is this chapter’s version of Chapter 2’s rule that legal is not the same as correct, and it has the same shape: the program runs, produces numbers, and means something other than what you intended.

What ModuleList is, and what it is not

nn.ModuleList is a module that registers the modules you put in it, using their positions as names. That is the whole of its job, and it is worth being precise about the two things people assume beyond it.

It does not define a computation:

layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
layers(torch.randn(1, 2))

NotImplementedError: Module [ModuleList] is missing the required "forward" function

Your forward still decides how, whether, in what order and how many times the contents are called. You can call them in reverse, call one twice, skip some based on the input, or never call one at all — and none of that changes what is registered.

That is the useful contrast with nn.Sequential, which also registers its children but additionally defines a forward computation that calls them in sequence.

For this simple example, both containers produce the same registered child names:

print("ModuleList:", [n for n, _ in nn.ModuleList([nn.Linear(2,2), nn.Linear(2,2)]).named_modules()])
print("Sequential:", [n for n, _ in nn.Sequential(nn.Linear(2,2), nn.Linear(2,2)).named_modules()])

ModuleList: ['', '0', '1']
Sequential: ['', '0', '1']

Same names, same registration, same checkpoint keys. Choose Sequential when the computation genuinely is “apply these in order”; choose ModuleList when forward needs to do something else with them. nn.ModuleDict, nn.ParameterList and nn.ParameterDict register the same way for the cases where you want string keys or bare parameters instead.

The second assumption worth dismantling is that ordinary Python containers are forbidden. They are not. What breaks is when a module is reachable only through an unregistered container. This is fine:

class Both(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = nn.Linear(2, 2)
        self.b = nn.Linear(2, 2)
        self.order = [self.a, self.b]      # plain list, contents registered by name

    def forward(self, x):
        for layer in self.order:
            x = layer(x)
        return x

['a.weight', 'a.bias', 'b.weight', 'b.bias']

Both layers are registered, because they were assigned to attributes. The list is a second, unregistered view of the same objects, used only for ordering. Nothing is lost. The rule is about reachability by the walk, not about which Python types you are allowed to use.

Buffers: model state that is not a parameter

Not everything a module owns should be trained. Running statistics, fixed masks, normalization constants, precomputed tables and counters can belong to the model without being trainable parameters. They may need to move with the module and may need to appear in its saved state, but they should not be discovered through model.parameters().

That is the relationship a registered buffer represents.

class Normalize(nn.Module):
    def __init__(self, features):
        super().__init__()
        self.register_buffer("mean", torch.zeros(features))
        self.scale = torch.ones(features)              # ordinary attribute
        self.gain = nn.Parameter(torch.ones(features))

    def forward(self, x):
        return (x - self.mean) * self.scale * self.gain

Three tensors, three relationships to the module:

named_parameters: ['gain']
named_buffers   : ['mean']
state_dict keys : ['gain', 'mean']

scale appears in none of them. It works perfectly in forward, because forward uses the Python object graph, and it is invisible to everything else.

The three-way distinction is worth holding precisely, because each row is a different set of answers rather than a different kind of tensor:

in parameters() in buffers() in state_dict() converted by .to() given to an optimizer
registered parameter yes no yes yes if you pass it
persistent buffer no yes yes yes not via model.parameters()
non-persistent buffer no yes no yes not via model.parameters()
ordinary tensor attribute no no no no no

Two entries in that table are commonly overstated, so state them narrowly. A buffer is not “a tensor that cannot have gradients” — it is state that is owned by the module and is not a parameter, and defining it through autograd behavior gets the concept wrong. And “converted by .to()” is doing more work than it looks; see below.

Why model.to(device) moves things, and what it leaves behind

model.to(...) does not search your object graph for tensors. It walks the registered parameters and buffers, exactly like named_parameters() does, and converts each one. Anything outside that walk stays where it is.

You do not need a GPU to see this. PyTorch has a meta device, which holds shape and dtype but no data — it exists for exactly this kind of structural question, and it is available on every machine:

m = Normalize(4)
m.to("meta")

print("gain :", m.gain.device)
print("mean :", m.mean.device)
print("scale:", m.scale.device)

m(torch.randn(2, 4, device="meta"))

gain : meta
mean : meta
scale: cpu
RuntimeError: Tensor on device cpu is not on the expected device meta!

The parameter moved. The buffer moved. The ordinary attribute did not, and the forward pass then failed on the mismatch. On a real GPU the same structure produces the message you have probably already met:

RuntimeError: Expected all tensors to be on the same device,
but found at least two devices, cuda:0 and cpu!

The instinctive repair is to add .to(x.device) at the failing line. That will make the error go away, and on the next line where scale is used it will come back, and it will keep coming back until every use site is patched. The structural question is better:

Is the wrong-device tensor outside the registered state tree, and if so, does it belong in it?

If it is genuine model state, register_buffer puts it inside the walk and the problem never recurs. If it is a temporary derived from the input, it should probably be created from the input in forward and not stored on the module at all.

Two details about to() on a module are worth knowing because they differ from to() on a tensor.

It mutates in place and returns itself. Chapter 2 established that tensor.to(...) returns a new tensor and leaves the original alone. A module converts the data held by each registered parameter and buffer without replacing the Parameter objects:

layer = nn.Linear(2, 2)
w = layer.weight
print("returns self:", layer.to(torch.float64) is layer)
print("same object :", layer.weight is w, "| dtype:", layer.weight.dtype)

returns self: True
same object : True | dtype: torch.float64

That object identity is load-bearing. An optimizer holds references to parameter objects, so an optimizer constructed before model.to("cuda") still holds the right ones afterwards. (Optimizer state — momentum buffers and the like — is created lazily on first use, so moving the model before the first step avoids any question about where that state lives.)

Explicit dtype conversion applies to floating-point and complex state, while integral state keeps its dtype. A .to(torch.float64) call does not turn your integer counters and boolean masks into floats:

layer.register_buffer("counter", torch.zeros(3, dtype=torch.long))
layer.to(torch.float64)
print(layer.weight.dtype, layer.counter.dtype)

torch.float64 torch.int64

Device movement applies to all of it; dtype conversion is selective. That is the behavior you want, and it is worth knowing it is deliberate rather than assuming every buffer follows every conversion.

Persistent and non-persistent buffers

Registration answers “does this belong to the module?” A second, independent question is “does this belong in a checkpoint?” For most state the answers match. Sometimes they do not, and persistent=False is how you say so.

The clearest case is state that is fully determined by the configuration:

class Block(nn.Module):
    def __init__(self, dim, max_len):
        super().__init__()
        self.proj = nn.Linear(dim, dim)
        self.register_buffer(
            "causal_mask",
            torch.tril(torch.ones(max_len, max_len, dtype=torch.bool)),
            persistent=False,
        )
        self.register_buffer("calls_seen", torch.zeros((), dtype=torch.long))

named_buffers   : ['causal_mask', 'calls_seen']
state_dict keys : ['calls_seen', 'proj.weight', 'proj.bias']

The mask is registered, so it moves to the GPU with everything else and you never write a device transfer for it. It is not in the checkpoint, because it contains no learned information: torch.tril(torch.ones(...)) reconstructs it exactly, every time, from max_len. Saving it would grow every checkpoint by a max_len × max_len array and, worse, would bake the training-time sequence length into the file, so that loading it into a model configured for a different length becomes a shape error over a value that could have been recomputed.

The counter is persistent, because how many batches the module has seen is not recoverable from the configuration. That distinction — recomputable from config versus accumulated from data — is usually the right test.

The observed ordering also reflects the recursive traversal in this example, but do not make correctness depend on state-dict ordering. The keys and their paths are the useful contract here.

state_dict() is the tree, keyed by path

We can now say precisely what a model state_dict contains. It does not serialize the model’s Python structure and it does not collect arbitrary Python attributes. It contains references to the module’s registered parameters and persistent buffers, keyed by their paths through the registered tree.

A complete training checkpoint may contain more than this — for example optimizer state, scheduler state, counters or metadata — but the model portion is usually represented by this state dictionary.

Everything people find mysterious about checkpoints follows from that sentence. Keys go missing when state stops being registered. Keys appear unexpectedly when state starts being registered. Keys get renamed when the path changes — and the path changes for reasons that look purely cosmetic in the source.

Here are two versions of a model, differing by one attribute name:

class V1(nn.Module):
    def __init__(self, w=16, d=2):
        super().__init__()
        self.stem = nn.Linear(2, w)
        self.blocks = nn.ModuleList([nn.Linear(w, w) for _ in range(d)])
        self.head = nn.Linear(w, 2)

class V2(nn.Module):                      # after a rename
    def __init__(self, w=16, d=2):
        super().__init__()
        self.stem = nn.Linear(2, w)
        self.encoder = nn.ModuleList([nn.Linear(w, w) for _ in range(d)])
        self.head = nn.Linear(w, 2)

Rather than loading and reading the traceback, compare the key sets directly. This is the same move as Chapter 2’s shape trace: line up expected against actual and find the first place they part company.

checkpoint = V1().state_dict()
model = V2()

have = list(checkpoint.keys())
want = list(model.state_dict().keys())

for a, b in zip(have, want):
    print(f"{a:24s} {b:24s} {'' if a == b else '<-'}")

print()
print("missing from checkpoint :", sorted(set(want) - set(have)))
print("unexpected in checkpoint:", sorted(set(have) - set(want)))

stem.weight              stem.weight              
stem.bias                stem.bias                
blocks.0.weight          encoder.0.weight         <-
blocks.0.bias            encoder.0.bias           <-
blocks.1.weight          encoder.1.weight         <-
blocks.1.bias            encoder.1.bias           <-
head.weight              head.weight              
head.bias                head.bias                

missing from checkpoint : ['encoder.0.bias', 'encoder.0.weight',
                           'encoder.1.bias', 'encoder.1.weight']
unexpected in checkpoint: ['blocks.0.bias', 'blocks.0.weight',
                           'blocks.1.bias', 'blocks.1.weight']

The two sets are the same size, the same shapes, in the same order, and they diverge at exactly one path component. load_state_dict reports the same thing in less structured form:

RuntimeError: Error(s) in loading state_dict for V2:
	Missing key(s) in state_dict: "encoder.0.weight", "encoder.0.bias", ...
	Unexpected key(s) in state_dict: "blocks.0.weight", "blocks.0.bias", ...

Reading it as a diff makes the fix obvious: rename the attribute back, or remap the keys, depending on which side is authoritative. It also makes the tempting alternative obviously wrong. strict=False does not repair anything; it downgrades the error to a return value:

print(model.load_state_dict(checkpoint, strict=False))

_IncompatibleKeys(missing_keys=['encoder.0.weight', 'encoder.0.bias', ...],
                  unexpected_keys=['blocks.0.weight', 'blocks.0.bias', ...])

The two encoder layers — four parameter tensors, weight and bias for each — remain at the new model’s initialization even though you may believe you loaded them from the checkpoint — which is precisely the failure this chapter opened with, arrived at by a different route. strict=False is a legitimate tool when you intend a partial load and can say which keys should be absent and why. Used to make a message go away, it converts a loud failure into a silent one.

One more property of state_dict() follows from the mechanism and catches people out. By default the tensors returned in a state_dict() are detached from autograd, but the state dictionary is a shallow snapshot of references to the module’s current state rather than cloned parameter data:

layer = nn.Linear(3, 2)
sd = layer.state_dict()
print("is the Parameter object:", sd["weight"] is layer.weight)
print("shares storage         :", sd["weight"].data_ptr() == layer.weight.data_ptr())

is the Parameter object: False
shares storage         : True

So the common early-stopping idiom does not do what it appears to:

best_state = model.state_dict()                  # shallow state references
best_state = copy.deepcopy(model.state_dict())   # independent in-memory snapshot

torch.save is unaffected, because it serializes immediately. Holding a state_dict in memory across further optimizer steps is where this bites: the “best” weights track the current ones, and the model you restore at the end is the last one, not the best one.

eval() is not no_grad()

Two settings are routinely confused, partly because they are usually used together. They control different mechanisms and are independent. The cleanest way to see it is to vary both and watch two different things change.

torch.manual_seed(0)
model = nn.Sequential(
    nn.Linear(8, 8),
    nn.Dropout(0.5),
    nn.Linear(8, 2),
)
x = torch.randn(4, 8)

def probe(label, out, m):
    grad_fn = type(out.grad_fn).__name__ if out.grad_fn is not None else "None"
    print(f"{label:22s} training={str(m.training):5s} "
          f"requires_grad={str(out.requires_grad):5s} "
          f"grad_fn={grad_fn:16s} out[0,0]={out[0, 0].item(): .4f}")

model.train()
torch.manual_seed(1); probe("train()", model(x), model)

model.eval()
torch.manual_seed(1); probe("eval()", model(x), model)

model.train()
with torch.no_grad():
    torch.manual_seed(1); probe("train() + no_grad()", model(x), model)

model.eval()
with torch.no_grad():
    torch.manual_seed(1); probe("eval() + no_grad()", model(x), model)

train()                training=True  requires_grad=True  grad_fn=AddmmBackward0   out[0,0]=-1.2312
eval()                 training=False requires_grad=True  grad_fn=AddmmBackward0   out[0,0]=-0.3193
train() + no_grad()    training=True  requires_grad=False grad_fn=None             out[0,0]=-1.2312
eval() + no_grad()     training=False requires_grad=False grad_fn=None             out[0,0]=-0.3193

Read it as two independent columns. eval() changed the output value and left requires_grad and grad_fn untouched. no_grad() changed requires_grad and grad_fn and left the output value untouched, to the last digit.

  • train() / eval() set a boolean called training on every registered module. Modules whose forward behavior depends on it — dropout, batch normalization, and anything you write that reads self.training — change what they compute. Autograd is not involved.
  • torch.no_grad() tells autograd not to record. Nothing about any module’s behavior changes; the same arithmetic runs and produces the same numbers, and there is simply no graph afterwards.

Calling model.eval() and expecting gradient tracking to stop is a real and common error. So is the reverse. Here is the reverse, using batch normalization, whose running statistics are buffers that get written to during a training-mode forward pass:

bn = nn.BatchNorm1d(4)
x = torch.randn(32, 4) * 3 + 5

print("initial          :", bn.running_mean)
with torch.no_grad():
    bn(x)
print("no_grad, train() :", bn.running_mean)

bn.eval()
with torch.no_grad():
    bn(x)
print("no_grad, eval()  :", bn.running_mean)

initial          : tensor([0., 0., 0., 0.])
no_grad, train() : tensor([0.4401, 0.5089, 0.5627, 0.5513])
no_grad, eval()  : tensor([0.4401, 0.5089, 0.5627, 0.5513])

torch.no_grad() did not protect the buffer. Running a validation pass on a model still in training mode mutates its running statistics with validation data — a silent corruption of model state that no gradient-related tool would catch, and one that shows up later as a train/eval discrepancy nobody can explain. eval() is what prevents it.

Since the flag propagates through the registered tree, it reaches exactly the modules the walk reaches, which gives the opening bug one final sting:

class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.registered = nn.Dropout(0.5)
        self.hidden = [nn.Dropout(0.5)]

m = M()
m.eval()
print("registered dropout .training:", m.registered.training)
print("hidden dropout .training    :", m.hidden[0].training)

registered dropout .training: False
hidden dropout .training    : True

A dropout layer stored only in a plain list remains in training mode when the registered model is switched to evaluation mode. Its validation outputs therefore remain stochastic and do not represent the evaluation behavior the programmer intended — while nothing in the surrounding training loop necessarily looks wrong.

Freezing, and the three sets that decide whether a parameter moves

Freezing is the clearest case of two properties that look like one:

model = nn.Sequential(nn.Linear(4, 4), nn.ReLU(), nn.Linear(4, 2))

for p in model[0].parameters():
    p.requires_grad_(False)

for name, p in model.named_parameters():
    print(f"{name:10s} requires_grad={p.requires_grad}")

print("state_dict keys:", list(model.state_dict().keys()))

F.cross_entropy(model(torch.randn(16, 4)), torch.randint(0, 2, (16,))).backward()

for name, p in model.named_parameters():
    print(f"{name:10s} grad={'None' if p.grad is None else 'present'}")

0.weight   requires_grad=False
0.bias     requires_grad=False
2.weight   requires_grad=True
2.bias     requires_grad=True
state_dict keys: ['0.weight', '0.bias', '2.weight', '2.bias']

0.weight   grad=None
0.bias     grad=None
2.weight   grad=present
2.bias     grad=present

The frozen parameters are still registered, still in named_parameters(), still in the checkpoint, and will still follow the model across devices. requires_grad_(False) changed their relationship to autograd without changing their registration. They remain model parameters, remain in the state dictionary and continue to move with the module, but ordinary backward computation no longer accumulates gradients into them. This is what you want from freezing a pretrained backbone — the weights are part of the model and must be saved and moved; they simply must not be updated.

For the conventional workflow in this book, there are three separate membership questions worth checking:

REGISTERED MODEL PARAMETERS model.parameters() What the module tree exposes for discovery.

AUTOGRAD-REACHED TENSORS Which tensors received gradients from this computation.

THIS OPTIMIZER’S PARAMETERS Which tensor objects this optimizer was actually given.

Do not turn that diagram into a mathematical rule that registration is required for an optimizer to modify a tensor. An optimizer can be handed a tensor directly even if it does not belong to any module.

For ordinary torch.optim training, the immediate update question is narrower:

  1. Is this exact tensor object in one of the optimizer’s parameter groups?
  2. What is its current .grad state?
  3. Given this optimizer’s rules and state, should a step change it?

A None gradient generally causes PyTorch optimizers to skip that tensor altogether, whereas an explicit zero gradient can behave differently because momentum, weight decay or other optimizer state may still matter.

Chapter 4 separated differentiation from update as two of its six questions. This is where the separation becomes mechanical: the optimizer is not a live view of the model. It is a list of tensor objects captured at construction time, and it will faithfully update those and nothing else.

x = torch.randn(32, 4)
y = torch.randint(0, 2, (32,))

model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))
optimizer = torch.optim.SGD(model[2].parameters(), lr=0.5)     # head only

before = {n: p.detach().clone() for n, p in model.named_parameters()}
for _ in range(20):
    loss = F.cross_entropy(model(x), y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

held = {id(p) for g in optimizer.param_groups for p in g["params"]}
for name, p in model.named_parameters():
    print(f"{name:10s} grad={'yes' if p.grad is not None else 'no':3s} "
          f"in_optimizer={str(id(p) in held):5s} "
          f"moved={(p - before[name]).abs().max().item():.6f}")

0.weight   grad=yes in_optimizer=False moved=0.000000
0.bias     grad=yes in_optimizer=False moved=0.000000
2.weight   grad=yes in_optimizer=True  moved=0.163894
2.bias     grad=yes in_optimizer=True  moved=0.200340

Registered, differentiated, gradient present, and completely stationary. If you only inspect gradients, this model looks healthy.

Passing a subset deliberately is a normal thing to do — fine-tuning a head, using different learning rates for different groups. The accidental version is easy to produce, because it can arise from a refactor that never touches the optimizer at all:

model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)

model[2] = nn.Linear(8, 2)      # swap the head after the optimizer exists

held = {id(p) for g in optimizer.param_groups for p in g["params"]}
for name, p in model.named_parameters():
    print(f"  {name:10s} in_optimizer={id(p) in held}")

  0.weight   in_optimizer=True
  0.bias     in_optimizer=True
  2.weight   in_optimizer=False
  2.bias     in_optimizer=False

The new head is correctly registered on the model. named_parameters() lists it, state_dict() saves it, .to() moves it. The optimizer still holds the old head’s parameter objects, which are now unreachable from the model.

That does not mean those stale objects will continue changing. Once the old head no longer participates in the forward pass, it normally receives no new gradient; with .grad is None, standard PyTorch optimizers skip it. The important failure is the opposite one: the new head receives gradients but is absent from the optimizer, so it never steps.

Both structures are internally consistent. They simply contain different parameter objects.

Membership is by object identity, not by name or by value, which is why id(p) in held is the check that answers the question. It is also why model.parameters() deduplicates shared tensors before you hand them over — an optimizer is entitled to complain about the same tensor appearing twice:

ValueError: some parameters appear in more than one parameter group

The module tree is not the forward graph

One last distinction, and it is the one that makes the tree comprehensible for real architectures. Registration says what the model owns. The forward pass says what it does. Neither constrains the other very much.

A module registered once can execute any number of times:

class Recurrentish(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.cell = nn.Linear(dim, dim)

    def forward(self, x, steps=3):
        for _ in range(steps):
            x = torch.tanh(self.cell(x))
        return x

modules: ['', 'cell']
params : ['cell.weight', 'cell.bias']
cell.weight.grad norm: 2.5847

One registered module, one weight, one bias — and three applications in the autograd graph, whose gradients accumulate into that single weight exactly as Chapter 3’s repeated-leaf example did. This is the mechanism behind recurrence, weight tying, and shared components in transformers. The registered tree is independent of how many times a module is called during a particular forward pass. The computation graph can contain one use, ten uses or no use of the same registered child.

The module tree itself is not immutable — assigning or replacing modules can change it, as the head-replacement example already showed. What matters is that executing forward() does not create another registered copy merely because a module was called again.

The reverse case is a module registered under two names:

class Shared(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.a = nn.Linear(dim, dim)
        self.b = self.a          # same object, second name

s = Shared(4)

print("named_modules   :", [n for n, _ in s.named_modules()])
print("named_parameters:", [n for n, _ in s.named_parameters()])
print("  remove_duplicate=False:",
      [n for n, _ in s.named_parameters(remove_duplicate=False)])
print("state_dict keys :", list(s.state_dict().keys()))
print("a.weight is b.weight:", s.a.weight is s.b.weight)

named_modules   : ['', 'a']
named_parameters: ['a.weight', 'a.bias']
  remove_duplicate=False: ['a.weight', 'a.bias', 'b.weight', 'b.bias']
state_dict keys : ['a.weight', 'a.bias', 'b.weight', 'b.bias']
a.weight is b.weight: True

Three different answers about the same object, and all three are correct for their purpose. named_modules() and named_parameters() deduplicate, because you want each tensor once when counting parameters or building an optimizer. state_dict() does not, because it is a mapping from every path to the state at that path; both keys are present and both refer to the same storage. Tied embedding and output layers in a language model produce exactly this, which is worth remembering the next time a checkpoint appears to contain a duplicate.

A structural audit

When a model does something you cannot explain, print what PyTorch has rather than reasoning about what the source implies. Two short functions cover most of it.

The first reports registered state, one row per tensor, with the columns that distinguish the sets we have been separating:

def registered_report(model, optimizer=None):
    held = set()
    if optimizer is not None:
        held = {id(p) for g in optimizer.param_groups for p in g["params"]}
    state = set(model.state_dict().keys())

    print(f"{'name':26s} {'kind':7s} {'device':7s} "
          f"{'req_grad':9s} {'grad':5s} {'state':6s} {'optim':5s}")

    for name, t in list(model.named_parameters()) + list(model.named_buffers()):
        kind = "param" if isinstance(t, nn.Parameter) else "buffer"
        print(f"{name:26s} {kind:7s} {str(t.device):7s} {str(t.requires_grad):9s} "
              f"{('yes' if t.grad is not None else 'no'):5s} "
              f"{('yes' if name in state else 'NO'):6s} "
              f"{('yes' if id(t) in held else 'no'):5s}")

The second audit looks through ordinary attributes for tensors and modules that are not already reachable through the registered tree. It is deliberately shallow — one level into common Python containers — so treat it as a practical diagnostic, not a complete traversal of arbitrary Python object graphs.

def unregistered_report(model):
    registered_modules = {id(m) for m in model.modules()}
    registered_tensors = {id(t) for t in model.parameters()}
    registered_tensors.update(id(t) for t in model.buffers())

    for mod_name, module in model.named_modules():
        prefix = mod_name + "." if mod_name else ""

        for attr, value in vars(module).items():
            if attr.startswith("_"):
                continue

            if isinstance(value, dict):
                candidates = [(f"{attr}[{k!r}]", v) for k, v in value.items()]
            elif isinstance(value, (list, tuple)):
                candidates = [(f"{attr}[{i}]", v) for i, v in enumerate(value)]
            else:
                candidates = [(attr, value)]

            for label, item in candidates:
                path = prefix + label

                if isinstance(item, nn.Module):
                    if id(item) not in registered_modules:
                        print(f"{path:26s} {type(item).__name__}")

                elif isinstance(item, torch.Tensor):
                    if id(item) not in registered_tensors:
                        print(f"{path:26s} {type(item).__name__}")

vars(module) is the instance __dict__ — the fourth destination from the assignment experiment. Skipping the underscore-prefixed keys removes _parameters, _buffers and _modules, so what remains is exactly the ordinary Python attributes, and we report the ones holding tensors or modules, including one level of list, tuple and dict.

Run both on a model carrying every mistake in this chapter at once:

class Classifier(nn.Module):
    def __init__(self, width=16, depth=3):
        super().__init__()
        self.stem = nn.Linear(2, width)
        self.encoder = Encoder(width, depth)     # plain-list version
        self.head = nn.Linear(width, 2)
        self.scale = torch.ones(width)
        self.register_buffer("calls", torch.zeros((), dtype=torch.long))

    def forward(self, x):
        self.calls += 1
        return self.head(self.encoder(torch.relu(self.stem(x)) * self.scale))

model = Classifier()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
F.cross_entropy(model(torch.randn(8, 2)), torch.randint(0, 2, (8,))).backward()

registered_report(model, optimizer)
print("\nreachable but unregistered:")
unregistered_report(model)

name                       kind    device  req_grad  grad  state  optim
stem.weight                param   cpu     True      yes   yes    yes  
stem.bias                  param   cpu     True      yes   yes    yes  
head.weight                param   cpu     True      yes   yes    yes  
head.bias                  param   cpu     True      yes   yes    yes  
calls                      buffer  cpu     False     no    yes    no   

reachable but unregistered:
scale                      Tensor
encoder.blocks[0]          Linear
encoder.blocks[1]          Linear
encoder.blocks[2]          Linear

Nine lines that would have ended the investigation at the beginning. The first table shows the registered state, whether each tensor currently has a gradient, whether it appears in the model state dictionary, and whether the optimizer holds that exact object.

Those facts let us reason about training, saving and movement without collapsing them into one question. The second shows four objects that participate in every forward pass and appear in none of it.

You do not have to use these exact functions. What matters is that when a model surprises you, the fastest path is usually a printout of named_modules(), named_parameters(), named_buffers() and state_dict() keys, plus a check of what the optimizer is actually holding — before adjusting a learning rate, adding a .to(device), or reaching for strict=False.

Using AI on a registration question

Registration bugs are a bad fit for the way people usually ask an assistant for help, because the model runs. There is no traceback to paste, the code looks idiomatic, and “my layer isn’t training” invites a plausible, confident answer about learning rates, initialization or normalization. Any of those might even improve the model slightly, which is worse than not helping.

The book’s rule applies with particular force here: ask for checkable claims and discriminating measurements before asking for rewritten code. The productive prompt names the structures and forbids the fix:

Here is a PyTorch model. One tensor participates in the forward pass but does
not appear to be training. Do not rewrite the model and do not suggest a fix yet.

Give me the smallest piece of inspection code that separately answers each of
these, and tell me what result would indicate a problem:

1. Can ordinary Python reach the tensor from the model object?
2. Is it an nn.Parameter, a registered buffer, or an ordinary attribute?
3. Does it appear in model.named_parameters() or model.named_buffers()?
4. Does it appear in model.state_dict(), and under what key?
5. Does it participate in the computation that produced this loss?
6. Does it have a .grad after backward(), and is that grad finite?
7. Is its exact object identity present in this optimizer's parameter groups?
8. Does its value change across one optimizer.step()?
9. Will model.to(device) convert it?

Then tell me which of those boundaries the evidence says is broken, and only
then propose the smallest change that would fix that specific boundary.

The numbered list is doing the real work. It forces the assistant to treat questions that sound like one question as nine, which is the same discipline this chapter has been applying by hand. The answers are also short enough to verify yourself, which matters, because an assistant is as capable as anyone else of asserting that a tensor is registered without checking.

The per-object version of that inspection is small enough to keep in your head. For any tensor you can name:

t = model.encoder.blocks[0].weight

print("type            :", type(t).__name__)
print("registered param:", [n for n, p in model.named_parameters() if p is t])
print("registered buffer:", [n for n, b in model.named_buffers() if b is t])
param_names = [
    n for n, p in model.named_parameters(remove_duplicate=False)
    if p is t
]
buffer_names = [
    n for n, b in model.named_buffers(remove_duplicate=False)
    if b is t
]
state_keys = set(model.state_dict())

print("registered param:", param_names)
print("registered buffer:", buffer_names)
print("in state_dict   :", [
    n for n in param_names + buffer_names
    if n in state_keys
])
print("requires_grad   :", t.requires_grad)
print("grad            :", "None" if t.grad is None
                           else round(t.grad.norm().item(), 4))
print("in optimizer    :", any(p is t for g in optimizer.param_groups
                               for p in g["params"]))
print("device          :", t.device)

type            : Parameter
registered param: []
registered buffer: []
in state_dict   : []
requires_grad   : True
grad            : 0.0363
in optimizer    : False
device          : cpu

Read the first two lines together and you have the whole chapter in miniature. It is an nn.Parameter. It is not a parameter of this model. Registration is a relationship, and the only way to know whether it holds is to ask the model rather than the tensor.

Notice that state-dict membership is derived from the registered path rather than from tensor object identity. A state_dict() contains detached state entries rather than the original Parameter objects, so the model’s registered names are the more useful bridge between ownership and serialized state.

When something seems invisible

Collected as a procedure, in the order that eliminates the most possibilities per step:

  1. Print the registered structure. named_modules(), named_parameters(), named_buffers(), and state_dict() keys. Compare against the structure you meant to build. Most registration bugs are visible here and nowhere else.
  2. Identify what kind of object it is. Registered child module, registered parameter, registered buffer, or ordinary Python attribute. If it is the last one, ask whether it should be, and which of the other three it should be instead.
  3. Look for the container. If a module is missing, check whether it is held only inside a plain list, tuple or dict. unregistered_report finds these; so does reading __init__ with the four destinations in mind.
  4. Separate the gradient question from the update question. A gradient means autograd reached the tensor. Check id(p) against the optimizer’s parameter groups before concluding anything about training.
  5. Check that the parameter actually moved. Snapshot with detach().clone(), step, and diff. This is the only direct evidence of an update, and it is Chapter 4’s fifth question.
  6. For a device error, ask about registration before adding a transfer. If the tensor is outside the registered tree, .to(device) was never going to move it, and patching one use site will not stop the next one.
  7. For a checkpoint error, diff the key sets. Find the first path component where they diverge. Reach for strict=False only when you can state which keys should be missing and why.
  8. For a train/eval discrepancy, check module.training on the modules that care. Print the flag rather than trusting that model.eval() reached everything; it reaches registered modules.

The habit underneath all eight:

Inspect the structure PyTorch sees instead of inferring it from the source code.

What you should now be able to answer

Here is a model of the kind you will be handed: adapters bolted onto a small encoder, a frozen normalization layer, a gate, a position buffer, and a head that someone re-initialized. It runs. Work through it before reading the answers.

class Adapter(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.down = nn.Linear(dim, dim // 4)
        self.up = nn.Linear(dim // 4, dim)

class Model(nn.Module):
    def __init__(self, dim=32, depth=2):
        super().__init__()
        self.embed = nn.Linear(8, dim)
        self.adapters = [Adapter(dim) for _ in range(depth)]
        self.gate = torch.ones(dim, requires_grad=True)
        self.register_buffer("positions", torch.arange(dim).float(),
                             persistent=False)
        self.norm = nn.LayerNorm(dim)
        self.head = nn.Linear(dim, 4)

    def forward(self, x):
        x = self.embed(x) + self.positions
        for a in self.adapters:
            x = x + a.up(torch.relu(a.down(x)))
        return self.head(self.norm(x) * self.gate)

model = Model()

for p in model.norm.parameters():
    p.requires_grad_(False)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

model.head = nn.Linear(32, 4)       # re-initialize the head

How many parameter tensors does model.parameters() return, and how many numbers is that? model.parameters() currently returns six tensors containing 484 numbers: embed.weight, embed.bias, norm.weight, norm.bias, head.weight, head.bias.

Across the model’s intended parameter-like state there are 1,620 numbers once the two adapters and the gate are included. But they are not all currently trainable: the 64 LayerNorm values have requires_grad=False.

So:

registered parameter values: 484 parameter-like values intended by the model: 1,620 values currently requiring gradients: 1,556

The two adapters account for 1,104 values and gate for another 32, and none of those are discoverable through model.parameters() in the broken version.

Why are the adapters missing? They are in a plain Python list, so they are ordinary attributes of Model and never enter _modules. forward reaches them through the list, which is why the computation is correct. The traversal that builds named_parameters() never sees them.

Why is gate missing? It is an ordinary tensor, not an nn.Parameter, so the assignment landed in the instance __dict__. requires_grad=True made it differentiable and did nothing to make it a parameter.

Which tensors receive gradients after backward()? The trainable embed and current head parameters, gate, and every adapter weight and bias. Not norm.weight or norm.bias, which were frozen. Note that the set receiving gradients and the set returned by parameters() overlap without either containing the other.

Which parameters actually change during training? Only embed.weight and embed.bias. Running ten steps and diffing gives:

embed.weight   req_grad=True  grad=yes optim=True  moved=0.010044
embed.bias     req_grad=True  grad=yes optim=True  moved=0.010110
norm.weight    req_grad=False grad=no  optim=True  moved=0.000000
norm.bias      req_grad=False grad=no  optim=True  moved=0.000000
head.weight    req_grad=True  grad=yes optim=False moved=0.000000
head.bias      req_grad=True  grad=yes optim=False moved=0.000000

Three different reasons for three different zeros. norm is in the optimizer and has no gradient, because it is frozen. head has a gradient and is not in the optimizer, because the optimizer holds the parameter objects of the head that was replaced. The adapters and gate are not in the table at all.

Why is head not in the optimizer when the optimizer was built from model.parameters()? Because the optimizer was built one line earlier. model.head = nn.Linear(32, 4) registered new parameter objects on the model while the optimizer continued to hold references to the old ones.

The new head participates in the forward pass and receives gradients but is absent from the optimizer. The old head remains in the optimizer but no longer participates in the computation, so its gradient is normally None and standard optimizers skip it.

The two structures now refer to different heads. Membership is by object identity, so the check is id(p) against optimizer.param_groups, not the parameter’s name.

What appears in state_dict()? embed.weight, embed.bias, norm.weight, norm.bias, head.weight, head.bias. The frozen norm parameters are present, because freezing does not affect registration. positions is absent, because it is a non-persistent buffer. The adapters and gate are absent, because they are not registered at all.

Is the missing positions key a bug? No, and it is the one absence here that is deliberate. torch.arange(dim) reconstructs it exactly from the configuration, so saving it would add nothing except a shape constraint on reload. Distinguishing this case from the adapters is the point: both are missing from the checkpoint, and only one of them was a decision.

What does model.to("cuda") move? The six registered parameters and the positions buffer. Not gate, and not any adapter weight. The first forward pass then fails on a device mismatch, and the failing line will be a.down(x) or the multiplication by gate depending on which comes first, neither of which is where the problem is.

What does model.eval() reach? Every registered module reachable from Model: the root Model, embed, norm, and the current head.

The adapters are absent because they are stored only in the plain Python list, so they keep whatever training/evaluation mode they already had.

Linear and LayerNorm do not change their computation here merely because their training flag changes, which is why this particular model shows no numerical symptom. Put Dropout or BatchNorm inside an unregistered adapter and the difference becomes observable immediately. Nothing in this model behaves differently in eval mode, so it costs nothing here — but the same structure with a nn.Dropout or nn.BatchNorm1d inside Adapter would silently keep training-mode behavior through every validation pass.

What is the smallest set of repairs? Wrap the adapters in nn.ModuleList, and make gate an nn.Parameter. Then rebuild the optimizer, or build it after the head assignment. The forward method needs no changes, and neither does the arithmetic. What changes is the description of ownership.

What evidence would prove the repair worked? Not that the loss went down. len(list(model.parameters())) becomes 15 and the parameter count becomes 1,620. state_dict() gains the adapter and gate keys and still omits positions. Every registered parameter except the frozen pair reports in_optimizer=True. A before/after diff shows the adapters, gate and head moving, and norm still at zero. And a save-then-reload round trip reproduces the same validation number rather than falling back toward chance.

Exercises

These are written to be run, and they map onto the notebook that accompanies this chapter.

  1. Predict the seven properties. For the BrokenLinear class, write down your prediction for requires_grad, is_leaf, grad after backward(), named_parameters(), state_dict(), dtype after .to(torch.float64), and whether an optimizer can update it. Then run it. Then change the single word that makes it nn.Parameter and rerun the identical probes. List which properties changed and which were already true.

  2. Break registration four different ways. Write one model with a plain-list container, one with a raw tensor used as a weight, one with a tensor attribute that should be a buffer, and one with a module assigned inside a plain dict. Run registered_report and unregistered_report on each. Confirm that the forward pass works in all four.

  3. The partial-training signature. Take the opening Classifier, train it, and produce a table of every layer’s maximum parameter change over training. Then argue from that table alone which layers were in the optimizer, without looking at the model source.

  4. Persistent or not. Add a causal_mask buffer and a steps_trained buffer to a small module. Decide which should be persistent and justify it by asking what would be lost on reload. Then verify with state_dict() keys, save and load, and check both buffers after loading.

  5. The 2×2. Reproduce the train()/eval() × no_grad() table for a model containing both dropout and batch normalization. Add a column for bn.running_mean before and after each forward pass, and state which of the four cells corrupts the running statistics.

  6. Three sets. Construct a model where one parameter is registered but frozen, one is registered and omitted from the optimizer, and one is registered, differentiated and updated. Predict grad, in_optimizer and moved for each before running registered_report.

  7. Diff a checkpoint. Save a model, then refactor it: rename a container, move a layer one level deeper, and wrap two layers in a Sequential. Without loading, predict the missing and unexpected key sets. Then write a function that remaps the old keys onto the new names, and prove it worked by comparing outputs on a fixed input.

  8. One module, many calls. Write a module that applies the same registered layer a number of times determined by an argument. Confirm the parameter count is independent of that number, and that the gradient magnitude is not. Then explain, from Chapter 3’s mechanism, why the gradients accumulate rather than overwrite.

  9. Audit something you did not write. Take any model from a library or a repository, run registered_report and unregistered_report on it, and answer: how many parameters, how many buffers, which buffers are non-persistent, and does anything appear in the second report? If something does, work out whether it is a bug or a deliberate choice.

Next: where does the input come from?

The model is now organized. A module is a tree of registered children, parameters and buffers; names are paths through that tree; and every framework operation you use daily is a traversal of it. parameters() gives the optimizer something to hold. state_dict() turns the tree into a checkpoint. .to() walks it to move state. train() and eval() walk it to set a flag. Autograd, separately, records what actually happened during a forward pass, and the two structures agree only to the extent that you made them agree.

That is enough structure to stop worrying about the size of a model. A hundred transformer blocks is the same walk as three linear layers, and the questions you ask when one of them misbehaves are the ones in this chapter.

Which leaves a gap on the other side. Every training loop in the book so far has had its data sitting in a tensor, already in memory, already the right shape, already on the right device. Real inputs arrive from disk, from a network, from files that need decoding, in formats that need converting, in quantities that do not fit in memory. Between the raw bytes and model(x) there is a pipeline, and it has to produce a batch before every single step of the training loop.

That pipeline turns out to have a failure mode this chapter’s tools cannot detect, because nothing about it is wrong. The model is correct, the gradients are correct, the checkpoints are correct — and the GPU spends most of its time doing nothing at all, waiting for the next batch to arrive.

Feeding the model.