Building Models with nn.Module
How nn.Module tracks parameters
Assigning a layer to self inside an nn.Module quietly adds its weights to a register, and everything else — training, saving, moving to GPU — reads from that register.
- 7 min read
- 3 reading levels
- Published
Read these first
On this page 5
One lesson, three depths. Pick the one that fits you today — you can switch any time.
Beginner — No maths. Plain English.
nn.Module is PyTorch's base class for models, and its main job is keeping a register of every trainable number inside your model.
Think of a school attendance register. When a new student joins the class, the teacher writes their name in it. From that day, the register decides everything — who sits the exam, who gets a report card, who is counted.
A PyTorch model works the same way. The moment you attach a layer to your model, its numbers go into the register. A layer that never made it into the register is invisible. It will not train, it will not be saved, and it will not move to the GPU.
Why it exists
A real model can hold millions of trainable numbers, spread across dozens of layers. Somebody has to keep the full list.
The optimizer needs the list, so it knows which numbers to nudge during training. Saving needs the list, so your trained model survives a restart. Moving to a GPU needs the list, so every number travels together.
Writing that list by hand, for every model, would be miserable and error-prone. So PyTorch made it automatic. You attach a part to your model, and the bookkeeping happens for you.
How it works
self.encoder = nn.Linear(4, 8)
|
v
[ nn.Module notices the assignment ]
|
v
register: encoder.weight (8 x 4 numbers)
encoder.bias (8 numbers)
|
v
optimizer, saving, GPU moves all read this registerThe register is why PyTorch code looks so clean. You describe the parts, and the plumbing is done for you.
A real example you have seen
Every model behind Google Photos search or your bank's fraud check was built this way. Someone described the parts, and the framework tracked the millions of numbers inside them.
When such a model is saved and shipped to your phone, the register is what gets written to disk.
Remember this
- Attaching a layer to your model registers its numbers automatically.
- The optimizer, saving, and GPU moves all read from that register.
- A layer hidden from the register silently does nothing — this is a real and common bug.
What to learn next
- Sequential, ModuleList and ModuleDict — the containers that register many layers at once.
- Buffers vs parameters — state that saves but never trains.
- PyTorch basics — if tensors themselves still feel new.
Developer — Code and libraries.
Setup
pip install torchWritten and tested against torch 2.5 on CPU. The behaviour here is core PyTorch and has been stable for years.
The register in action
import torch
from torch import nn
class RipenessNet(nn.Module):
def __init__(self):
super().__init__()
self.encoder = nn.Linear(4, 8) # assigning to self registers the layer
self.decoder = nn.Linear(8, 1)
self.scale = nn.Parameter(torch.ones(1))
self.extras = [nn.Linear(8, 8)] # hidden inside a plain list: invisible
def forward(self, x):
h = torch.relu(self.encoder(x))
h = torch.relu(self.extras[0](h))
return self.decoder(h) * self.scale
model = RipenessNet()
for name, p in model.named_parameters():
print(f"{name:16s} {tuple(p.shape)}")
print("registered:", sum(p.numel() for p in model.parameters()))
print("mentions 'extras':", any("extras" in n for n, _ in model.named_parameters()))
print("in state_dict:", any("extras" in k for k in model.state_dict()))scale (1,) encoder.weight (8, 4) encoder.bias (8,) decoder.weight (1, 8) decoder.bias (1,) registered: 50 mentions 'extras': False in state_dict: False
Read that output carefully
The layer inside the plain Python list is gone. It runs in forward, so the model produces output and nothing crashes. But its 72 numbers are not registered. The optimizer will never update them, and saving will never store them.
This is the worst kind of bug: the silent kind. Your model trains, the loss even falls a little, and one whole layer stays random forever.
What makes registration happen
Three assignments register things, and nothing else does:
| You assign | What gets registered |
|---|---|
an nn.Module (like nn.Linear) | the submodule and all its parameters |
an nn.Parameter | one trainable tensor |
an nn.ModuleList / nn.ModuleDict | every module inside it |
A plain list, dict, or bare tensor is stored as an ordinary attribute. PyTorch does not look inside it. The fix for the broken code above is one word: nn.ModuleList([nn.Linear(8, 8)]). That container is the subject of the next lesson.
super().__init__() must run before any layer assignment. It creates the internal register. Forget it and PyTorch raises AttributeError: cannot assign module before Module.__init__() call, which is the framework telling you the register does not exist yet.
Why scale printed first
Directly-attached parameters are listed before submodules. The order is a detail of the internal bookkeeping, not something your code should depend on. Depend on the names, never the order.
Common mistakes
Layers in a plain list or dict. Shown above. Symptom: the model underperforms and nothing errors. Check with sum(p.numel() for p in model.parameters()) against a hand count.
Creating layers inside forward. A layer built during forward is born fresh on every call, with random weights each time. Build layers in __init__, use them in forward.
A bare tensor where a parameter was meant. self.scale = torch.ones(1) stores a constant. self.scale = nn.Parameter(torch.ones(1)) stores a trainable number. Only the second appears in model.parameters().
Passing the wrong thing to the optimizer. torch.optim.SGD(model.parameters(), lr=0.1) is correct. Passing a fresh model's parameters, or a copied list from before a layer swap, trains the wrong tensors.
Try it yourself
Fix RipenessNet by wrapping extras in nn.ModuleList. Rerun and confirm the count rises from 50 to 122, and that extras.0.weight appears in the state dict.
What to learn next
- Sequential, ModuleList and ModuleDict — the containers that register many layers at once.
- Buffers vs parameters — state that saves but never trains.
- PyTorch basics — if tensors themselves still feel new.
Researcher — Mathematics and papers.
The mechanism is __setattr__
nn.Module overrides Python's __setattr__. Every assignment to self is inspected and routed into one of three internal dicts:
_parameters: dict[str, Parameter]— directnn.Parameterassignments._modules: dict[str, Module]— submodule assignments._buffers: dict[str, Tensor]— tensors added viaregister_buffer(see buffers vs parameters).
Anything else lands in the ordinary __dict__. named_parameters() is a recursive walk: yield _parameters, then recurse into _modules, prefixing names with the attribute path and a dot. That walk defines the state_dict key namespace — encoder.weight is literally "attribute encoder, then its attribute weight".
Deduplication matters: the walk tracks tensor identity, so a parameter shared between two submodules (weight tying) is yielded once. named_parameters(remove_duplicate=False) disables this.
Consequences of the design
Registration is orthogonal to execution. forward is ordinary Python; the graph is rebuilt on every call (define-by-run, Paszke et al., 2019, PyTorch: An Imperative Style, High-Performance Deep Learning Library, NeurIPS). Nothing checks that registered modules are used, or that used modules are registered. Both mismatches are legal and both are bugs in practice: the dead registered layer wastes memory and pollutes the state dict, and the unregistered live layer never trains.
model.to(device) works by walking the same register. Each Parameter object is kept and its storage is replaced, so references held elsewhere — an already-constructed optimizer, for instance — keep pointing at live tensors. Buffers, by contrast, are reassigned to new tensors, an asymmetry that captured references can observe.
requires_grad is per-tensor, registration is per-name. Freezing (p.requires_grad_(False)) leaves the parameter registered — it stays in state_dict and moves with .to(). Registration answers "does this belong to the model"; requires_grad answers "does autograd track it". Transfer learning leans on that distinction.
Cost
The register itself is negligible: a few dicts per module. The recursive walk is O(number of modules + parameters) and is re-run on every call to parameters(), state_dict(), and .to(). For models with tens of thousands of modules (large mixture-of-experts trees), this walk shows up in profiles, which is one motivation for flattened representations in torch.compile and FSDP's parameter grouping.
Reading
- Paszke et al. (2019), PyTorch: An Imperative Style, High-Performance Deep Learning Library — the design rationale, including define-by-run.
- The
nn.Modulesource is short and readable:torch/nn/modules/module.py, in particular__setattr__and_named_members.
What to learn next
- Sequential, ModuleList and ModuleDict — the containers that register many layers at once.
- Buffers vs parameters — state that saves but never trains.
- PyTorch basics — if tensors themselves still feel new.