Modules

Assign a nn.Linear to self.up inside a module's __init__ and you haven't just stored a reference the way a plain Python attribute would. nn.Module overrides attribute assignment to notice what you handed it, a parameter, a buffer, or another module, and registers that object into the tracking structures the module keeps for exactly this purpose.

Two categories cover everything a module can hold. Parameters are leaf tensors with requires_grad=True, the numbers a training loop is meant to update. Buffers are state without gradients: a running mean in batch normalization, something the module needs to carry forward but never wants a gradient computed for. parameters() walks every registered parameter across a module and its submodules; state_dict() does the same walk but returns names mapped to tensors instead of a bare iterator, covering both parameters and buffers together.

In PyTorch, state is the model.

import torch
from torch import nn

class Mlp(nn.Module):
    def __init__(self):
        super().__init__()
        self.up = nn.Linear(8, 32)
        self.down = nn.Linear(32, 1)

    def forward(self, x):
        return self.down(torch.relu(self.up(x)))

model = Mlp()
print(sum(p.numel() for p in model.parameters()))   # 321
print(list(model.state_dict())[:2])                 # ['up.weight', 'up.bias']

This is the philosophical opposite of the other path's answer to the same question. Over there, state is never something an object holds; it's an argument a function takes and a new value that function returns, and nothing is ever mutated in place. Here, a module's parameters live on self, get mutated by the optimizer step in the next chapter, and never appear in the function signature of forward at all.

Neither answer is a workaround for the other's limitation; each is chosen on purpose, for a reason that shows up two chapters from now. torch.compile and torch.export exist specifically to pull a static graph back out of this mutable, stateful world, one snapshot at a time. Recovering a graph from mutation is real work; the other path never has to do that work, because it never left graph form to begin with.

state_dict() is not a debugging tool you reach for occasionally. It's the actual serialization contract PyTorch checkpoints on: a flat mapping from dotted names, "up.weight", "down.bias", to tensors, and load_state_dict() reads that same mapping back into an existing module by matching names. Save a checkpoint, change nothing about the architecture, and state_dict() round-trips exactly; change a layer's name and loading fails loudly, on a name it can't find, rather than silently loading the wrong tensor into the wrong place.

Chapter 4's checkpoint is built entirely on top of this: the model's half of a saved checkpoint is just this dictionary, written to disk and read back.

Call model.train() or model.eval() and you haven't passed anything into forward. You've flipped a flag on the module itself, and every submodule that cares, dropout, batch normalization, checks that flag to decide how it behaves: dropout zeroes activations during training and passes everything through during eval; batch norm uses batch statistics while training and the running statistics it accumulated once it's in eval.

That's the same pattern as the first section, playing out at a smaller scale. The behavior a layer shows for the exact same input depends on state sitting on the object, not on anything you supplied to the call. Move a model from train to eval without noticing, and every dropout layer in it silently starts doing something different than it was a line before, for a reason nothing in your function call reveals.