Building models with nn.Module: Question bank

Before we start, a few words we lean on the whole way. Read these once so nothing below uses a term you haven't met.

Definition The vocabulary that recurs

Parameter — a tensor PyTorch should learn (update during training). Wrapped as nn.Parameter. Think: a knob that the optimizer is allowed to turn.
Optimizer — the object that actually turns the knobs. During training it looks at each parameter's gradient (the direction that reduces error) and nudges the parameter a little that way. In code it is created like optimizer = torch.optim.SGD(model.parameters(), lr=0.01) — notice it is handed exactly the list from model.parameters(), which is why unregistered tensors never get trained.
Registration — the act of PyTorch noticing a parameter or submodule so that .parameters(), .to(device), and saving all find it. Assigning self.x = nn.Parameter(...) inside __init__ registers it; putting it in a plain Python list does NOT.
Buffer — a tensor that is saved and moved with the model but never trained. Registered with self.register_buffer("name", tensor). Running averages inside BatchNorm are buffers: the model needs them at test time, but the optimizer must not touch them.
forward — the method where you describe the computation. You never call it directly; calling model(x) runs it for you.
nn.ModuleList — a list-like container that registers every module you put in it. Written self.layers = nn.ModuleList([nn.Linear(4,4), nn.Linear(4,4)]). It behaves like a Python list in a for loop but, unlike a plain list, its contents are tracked for training.
==@== — Python's matrix-multiplication operator. x @ w multiplies the matrix x by the matrix w, the same thing a linear layer does under the hood: rows of x combine with columns of w.

The parent note is Building models with nn.Module — everything here goes one level deeper into the traps that note's clean code hides.

Three pictures anchor the ideas the traps keep circling back to.

The registration tree — why .parameters(), .to(device), and the optimizer all "just work":

Forward-pass control flow — why you call model(x), never model.forward(x):

The residual skip — why the +identity path keeps gradients alive:

True or false — justify

TF1. Storing layers in a plain Python list self.layers = [nn.Linear(...), ...] trains them just like storing them individually.

False. A plain list is invisible to registration, so those layers' parameters never appear in .parameters() and the optimizer (the object that nudges parameters) never receives them. Use nn.ModuleList, which registers each contained module.

TF2. nn.ReLU() and the functional F.relu(x) produce identical numbers on the same input.

True. ReLU is stateless (no parameters), so the module and the function compute the exact same max(0, x). The module form only adds convenience (it appears in .modules() and prints in the model summary).

TF3. Forgetting super().__init__() still works as long as you define forward.

False. That call sets up the internal __setattr__ machinery. Without it, assigning self.fc1 = nn.Linear(...) raises an error (or silently fails to register), because the module's parameter/submodule dictionaries were never created.

TF4. Every tensor you create inside __init__ becomes a trainable parameter.

False. Only tensors wrapped in nn.Parameter are registered as parameters. A raw torch.randn(10) assigned to self.x is stored but is not trainable and won't move with .to(device) unless you register it as a buffer (state that is saved and moved but never trained).

TF5. model.parameters() returns only the parameters defined directly in that class, not those of its submodules.

False. It walks the submodule tree recursively, collecting its own parameters plus every submodule's. That recursion is exactly why composing modules "just works."

TF6. Calling model(x) and model.forward(x) are completely interchangeable.

Mostly false in practice. model(x) goes through __call__, which runs registered hooks and bookkeeping around forward. model.forward(x) skips all of that. Always use model(x).

TF7. nn.Parameter tensors have requires_grad=True by default.

True. That is the point of wrapping a tensor as a Parameter — it declares "track gradients for me." A plain tensor defaults to requires_grad=False.

TF8. Two modules that share the same nn.Parameter object will have that parameter counted twice by sum(p.numel() for p in model.parameters()) if you're not careful.

True (as a trap). Naive counting can double-count shared weights. Using model.parameters() at the top level de-duplicates by identity, but iterating submodules separately and summing can over-count tied weights.

Spot the error

SE1. self.layers = [nn.Linear(4,4) for _ in range(3)] then later for l in self.layers: x = l(x).

The forward pass runs, but the three layers are in a plain list → unregistered → never trained. Fix: wrap in nn.ModuleList([...]), the list-like container that registers each module.

SE2. def __init__(self): self.fc = nn.Linear(4,2) — no super().__init__() line.

Missing the parent constructor. Assigning self.fc before the module internals exist throws AttributeError: cannot assign module before Module.__init__() call.

SE3. self.w = torch.randn(4, 2) used as a weight inside forward with x @ self.w (where @ is Python's matrix-multiply operator).

torch.randn gives a plain tensor: not registered, requires_grad=False, won't train, won't move with .to(device). Fix: self.w = nn.Parameter(torch.randn(4, 2)).

SE4. def forward(self, x): return self.fc1(x); self.fc2(x).

The return on the first statement ends the function; self.fc2 never runs. The second layer's output is silently discarded. Chain them: x = self.fc1(x); return self.fc2(x).

SE5. A binary classifier ends with x = torch.sigmoid(self.fc3(x)) and is trained with nn.BCEWithLogitsLoss.

BCEWithLogitsLoss applies its own sigmoid internally. Applying sigmoid first squashes twice, wrecking the gradient. Either output raw logits with BCEWithLogitsLoss, or use plain BCELoss with an explicit sigmoid — never both.

SE6. Residual block writes out = self.relu(out); out = out + identity (ReLU before adding the skip).

The classic ResNet block adds the skip before the final activation: out = out + identity; out = self.relu(out). Activating the residual branch alone can clip the identity path's contribution and hurts gradient flow.

SE7. self.hidden = nn.ModuleList([...]) but in forward the loop is for layer in self.hidden_layers.

Attribute name mismatch — self.hidden vs self.hidden_layers — raises AttributeError. A registration name and its usage must match exactly.

SE8. Model defined with nn.Dropout but during evaluation the code never calls model.eval().

Dropout stays in training mode, randomly zeroing activations at test time → noisy, non-reproducible predictions. Call model.eval() (and model.train() to switch back) so stateful modules behave correctly.

Why questions

WHY1. Why subclass nn.Module at all instead of writing a function layer(x, w, b)?

A function forces you to pass every weight explicitly, tracks no gradients automatically, can't save/load, and can't switch train/eval. The Module owns its parameters and gives all of that for free.

WHY2. Why is the output layer of a classifier usually left with no activation?

The loss (e.g. BCEWithLogitsLoss, CrossEntropyLoss) expects raw logits and applies the softmax/sigmoid internally for numerical stability. Adding your own activation double-counts it.

WHY3. Why does the "+1" (the identity path) in a residual block fight vanishing gradients?

With a skip,

\partial L / \partial x = \partial L / \partial out \cdot (\partial F / \partial x + 1)

. Even if

\partial F / \partial x \approx 0

, the

+ 1

guarantees the gradient still reaches earlier layers instead of dying.

WHY4. Why store even a stateless ReLU as self.relu = nn.ReLU() rather than calling F.relu?

Consistency and introspection: it shows up in the module tree and print-out, and the pattern is identical to stateful modules (Dropout, BatchNorm) where storing-as-attribute is mandatory for train/eval switching.

WHY5. Why does .to(device) move everything with one call?

Registration builds a recursive tree of submodules and parameters. .to(device) walks that tree and moves each registered tensor — which is exactly why unregistered raw tensors get left behind on the wrong device.

WHY6. Why must you register a running statistic (like a mean you compute yourself) as a buffer rather than a parameter?

A buffer is state that should be saved and moved with the model but not trained by the optimizer. Making it a Parameter would let the optimizer overwrite it with gradient steps, corrupting the statistic.

WHY7. Why does nn.Linear beat a hand-written torch.randn weight matrix?

nn.Linear initializes weights with a principled scheme (Kaiming/Xavier) that keeps activation variance sane across layers, handles the bias correctly, and registers everything. Raw randn gives poorly scaled, unregistered weights.

Edge cases

EC1. What does model.parameters() return for a module with no nn.Parameter and no submodules (e.g. a pure nn.ReLU)?

An empty iterator — zero parameters. Stateless modules are perfectly valid; they simply contribute nothing to the trainable set.

EC2. You create nn.Parameter(torch.randn(0)) (a zero-length tensor). Is it a valid parameter?

Yes, it registers, but it holds no elements (numel() == 0), so it contributes nothing to training or the parameter count. Usually a sign of a logic bug in your dimension arithmetic.

EC3. What happens if two layers share the same weight object (weight tying) — how many parameters does the optimizer update?

One. Both usages point to the same tensor, so a single gradient (accumulated from both uses) updates it once. Top-level .parameters() de-duplicates it, but manual per-submodule counting will over-count.

EC4. num_layers = 0 is passed to a network that builds nn.ModuleList([... for _ in range(num_layers)]). Does it break?

No — you get an empty ModuleList. The forward loop simply never executes, so input flows straight from the input layer to the output layer. Degenerate but valid.

EC5. You reassign self.fc1 = nn.Linear(4, 4) a second time later in __init__. Which one is registered?

Only the last assignment. The re-assignment replaces the earlier registration under the same name; the first Linear's parameters vanish from the tree.

EC6. Input batch of size 1 vs a bare unbatched vector — does a nn.Linear-based model care?

nn.Linear operates on the last dimension, so a shape [1, in] works fine. A bare [in] also works for Linear, but BatchNorm and many modules require a batch dimension and error on unbatched input — always keep the leading batch axis.

EC7. During model.eval() with no Dropout/BatchNorm anywhere, does anything actually change?

No numerical difference — eval mode only alters stateful modules. Still, calling it is good hygiene so behavior is correct the moment you later add a Dropout or BatchNorm layer.

Recall Quick self-test

Why won't a plain Python list of layers train? ::: Its modules are never registered, so .parameters() never sees them and the optimizer never updates them. What does the +identity in a residual block guarantee for gradients? ::: A +1 term in $\partial L / \partial x$ , so gradients survive even when $\partial F / \partial x \to 0$ . Why leave classifier outputs as logits? ::: The loss applies sigmoid/softmax internally for numerical stability; adding your own double-applies it.