Building models with nn.Module: Worked examples

This page is the practice arena for Building models with nn.Module. The parent note taught you the pattern; here we hit every scenario that pattern can throw at you — from a plain layer, through the sneaky "my parameters vanished" bug, all the way to a degenerate zero-layer network.

Before we start, one word we lean on constantly:

The scenario matrix

Every mistake, exam twist, and real bug with nn.Module falls into one of these cells. Each worked example below is tagged with the cell(s) it covers.

Cell	Scenario class	What can go wrong / what's tested
A	Single `nn.Linear` — count parameters	off-by-bias errors, weight-shape confusion
B	Multi-layer stack, count total params	forgetting a layer, wrong hidden dim
C	Zero / degenerate: `num_layers = 0`	does the loop body ever run?
D	The bug cell: Python `list` vs `nn.ModuleList`	parameters silently NOT registered
E	Skip / residual dimension match	shapes must line up before `+`
F	train vs eval mode (Dropout/BatchNorm)	forgetting `.eval()`, non-determinism
G	Real-world word problem	translate a spec into a module
H	Exam twist: shared weights / `nn.Parameter`	one tensor counted once, tied layers

Below, we cover all eight cells across eight examples.

Example 1 — Cell A: a single Linear layer

Forecast: guess a number before reading on. Is it $15$ ? $18$ ? $8$ ?

Steps.

Write what nn.Linear computes. It computes $y = x W^{⊤} + b$ .
- Why this step? You cannot count parameters until you know which tensors exist. nn.Linear owns exactly two: the weight $W$ and the bias $b$ .
Shape of $W$ . nn.Linear(in, out) stores weight of shape $(out, in) = (3, 5)$ .
- Why this step? PyTorch stores the weight transposed so that the multiply is $x W^{⊤}$ . That is $3 \times 5 = 15$ numbers.
Shape of $b$ . The bias has one number per output: shape $(3,)$ → $3$ numbers.
- Why this step? Each output neuron gets its own additive shift, so bias length = out_features.
Total. $15 + 3 = 18$ learnable parameters.

$params = weight out \times in + bias out = 3 \times 5 + 3 = 18$

Verify: sum(p.numel() for p in layer.parameters()) returns $18$ ; layer.weight.shape == (3,5), layer.bias.shape == (3,). ✅ Units check: each parameter is one real number, no double counting.

Example 2 — Cell B: a full stack (parent's Example 1)

Forecast: the parent claimed 117121. Try to reproduce it before checking.

Steps.

fc1: $784 \to 128$ . Weight $128 \times 784 = 100352$ , bias $128$ . Subtotal $100480$ .
- Why this step? Use the Cell-A rule per layer: $out \times in + out$ .
fc2: $128 \to 128$ . Weight $128 \times 128 = 16384$ , bias $128$ . Subtotal $16512$ .
- Why this step? Same rule; here in = out = 128.
fc3: $128 \to 1$ . Weight $1 \times 128 = 128$ , bias $1$ . Subtotal $129$ .
- Why this step? Output layer for binary classification produces a single logit.
Add. $100480 + 16512 + 129 = 117121$ .

Verify: sum(p.numel() for p in SimpleClassifier(784,128,1).parameters()) == 117121. ✅

Example 3 — Cell C: the degenerate zero-layer case

Forecast: crash? empty list? Guess.

Steps.

Look at the loop. nn.ModuleList([nn.Linear(20,20) for _ in range(0)]) — the comprehension range(0) produces nothing.
- Why this step? Degenerate inputs live in the loop bounds. range(0) is empty, so zero hidden Linears are created. No crash: an empty nn.ModuleList is perfectly valid.
What survives? Only input_layer ( $10 \to 20$ ) and output_layer ( $20 \to 4$ ).
- Why this step? The network still transforms input to output — it just has no middle.
Count. input_layer: $20 \times 10 + 20 = 220$ . output_layer: $4 \times 20 + 4 = 84$ . Total $= 304$ .
- Why this step? Cell-A rule per surviving layer.
Forward still runs. The for layer in self.hidden_layers: loop body simply never executes; data flows input_layer → output_layer.

Verify: with num_layers=0, len(model.hidden_layers) == 0 and total params $= 220 + 84 = 304$ . ✅ Sanity: adding one hidden layer would add $20 \times 20 + 20 = 420$ more.

Example 4 — Cell D: the bug that eats your parameters

Worked example Why did my model not train?

A beginner writes the hidden layers as a plain Python list:

self.hidden = [nn.Linear(8, 8) for _ in range(3)]

instead of nn.ModuleList([...]). Each nn.Linear(8,8) has $8 \times 8 + 8 = 72$ params, so $3$ of them = $216$ params. How many of those $216$ does model.parameters() actually see?

Forecast: $216$ ? $72$ ? $0$ ? Commit to an answer.

Steps.

Recall the clipboard rule. From the definition above: a plain list is not on the clipboard. PyTorch's __setattr__ magic only registers nn.Module, nn.Parameter, and official containers.
- Why this step? Registration is what makes .parameters() find something. No registration → invisible.
Trace what happens. self.hidden = [...] stores a normal list. PyTorch sees "a list" and shrugs — it does not peek inside to find the nn.Linear objects.
- Why this step? Look at the figure: the three orange layers sit outside the clipboard, so the collector's recursive walk never reaches them.
Count visible params. If the rest of the model has no other layers, sum(p.numel() for p in model.parameters()) == 0.
- Why this step? Zero registered parameters. The optimizer gets an empty list; training does literally nothing.
The fix. Wrap it: self.hidden = nn.ModuleList([...]). Now all $216$ appear.
- Why this step? nn.ModuleList is on the clipboard, so its contents are registered recursively.

Verify: plain-list version → $0$ registered params; nn.ModuleList version → $216$ ( $= 3 \times 72$ ). ✅

Example 5 — Cell E: residual shapes must match

Forecast: does the + work, or throw a shape error?

Steps.

Spatial size under padding=1, kernel=3. Output size $= H - k + 2 p + 1 = 8 - 3 + 2 (1) + 1 = 8$ . Same for $W$ .
- Why this step? You may only add tensors of identical shape. So we must check the conv preserves $H, W$ . With $k = 3, p = 1$ it does — that is why residual blocks use this exact padding.
Channels. Conv2d(64, 64, ...) keeps channels at $64$ .
- Why this step? The addition also needs matching channel count. In-channels = out-channels = 64 by design.
Compare shapes. $F (x)$ is $[32, 64, 8, 8]$ ; $x$ is $[32, 64, 8, 8]$ . Identical → the + is legal, element-wise.
- Why this step? Look at the figure: both boxes are the same size, so the arrows line up one-to-one.
Result shape. $[32, 64, 8, 8]$ , unchanged.

Verify: conv output spatial $= 8 - 3 + 2 \cdot 1 + 1 = 8$ ; shapes $[32, 64, 8, 8]$ equal; sum shape $[32, 64, 8, 8]$ . ✅

Example 6 — Cell F: train vs eval mode

Forecast: which mode gives repeatable outputs — train or eval?

Steps.

What Dropout does in training. With model.train(), Dropout randomly zeros each element with probability $p = 0.5$ , then scales the survivors by $\frac{1}{1 - p} = 2$ .
- Why this step? The random mask changes every forward pass, so two calls on the same input give different outputs. This regularizes.
Why the $\times 2$ scale. Expected value must be preserved: $E [output] = (1 - p) \cdot \frac{1}{1 - p} \cdot x = x$ .
- Why this step? So that switching to eval doesn't shift the average activation.
What Dropout does in eval. With model.eval(), Dropout becomes the identity — no zeroing, no scaling. Output = input, exactly, every time.
- Why this step? At test time you want deterministic, full-strength predictions. The mode flag flips this behaviour.
Answer. train → outputs differ (random mask). eval → outputs identical (pass-through).

Verify: in eval mode, dropout(x) equals x; expected scaling factor $\frac{1}{1 - 0.5} = 2$ . ✅

Example 7 — Cell G: real-world word problem

Forecast: roughly how many parameters — tens of thousands, hundreds of thousands, millions?

Steps.

Translate the words into dimensions. Input $= 28 \times 28 = 784$ . Hidden $= 256$ . Output $= 10$ .
- Why this step? "Flattened 28×28" is the literal input length; "10 class scores" is the output length. Everything else is architecture.

Write it.

class MNISTNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 256)
        self.fc3 = nn.Linear(256, 10)
        self.relu = nn.ReLU()
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        return self.fc3(x)   # raw logits, no activation

Why this step? Three Linear layers realize the two-hidden-layer spec; ReLU between them; the final layer stays "bare" because CrossEntropyLoss expects logits, not probabilities.

Count.
- fc1: $256 \times 784 + 256 = 200704 + 256 = 200960$ .
- fc2: $256 \times 256 + 256 = 65536 + 256 = 65792$ .
- fc3: $10 \times 256 + 10 = 2560 + 10 = 2570$ .
- Why this step? Cell-A rule per layer.
Total. $200960 + 65792 + 2570 = 269322$ .

Verify: total params $= 200960 + 65792 + 2570 = 269322$ . ✅ Sanity: hundreds of thousands, as forecast for a modest MLP.

Example 8 — Cell H: the exam twist (weight tying)

Worked example Shared weights — how many parameters?

An encoder–decoder ties its two Linear layers: the decoder reuses the encoder's weight matrix (this is real — it's used in language models to tie input and output embeddings).

class Tied(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.enc = nn.Linear(d, d, bias=False)
        self.dec = nn.Linear(d, d, bias=False)
        self.dec.weight = self.enc.weight   # tie!
    def forward(self, x):
        return self.dec(self.enc(x))

With d = 4, how many unique learnable parameters does model.parameters() return? (Naively two $4 \times 4$ matrices = 32.)

Forecast: $32$ (two matrices) or $16$ (one shared matrix)?

Steps.

Count naively. Two nn.Linear(4,4, bias=False) → each weight $4 \times 4 = 16$ , total would be $32$ .
- Why this step? This is the trap answer. We must check whether the tie changes the count.
The tie makes them one object. self.dec.weight = self.enc.weight sets the decoder's weight to be the exact same tensor object as the encoder's.
- Why this step? Look at the figure: both layers point at a single shared box of $16$ numbers. There is one tensor, referenced twice.
.parameters() de-duplicates. PyTorch's parameter iterator returns each unique tensor once, even if several modules point to it.
- Why this step? Counting it twice would apply the gradient twice; de-duplication keeps a single update per tensor.
Answer. $16$ unique parameters, not $32$ .

Recall Why tie weights at all?

Tying enc/dec weights ::: it halves parameters and, in language models, forces the input and output embedding spaces to agree — often improving generalization.

Verify: len({id(p) for p in model.parameters()}) gives $1$ unique tensor of $16$ numbers → total numel $= 16$ . ✅

Recall Scenario matrix — self test

Plain Python list of layers: how many params does .parameters() see? ::: Zero — a list is not registered; use nn.ModuleList. DeepNetwork(..., num_layers=0, ...): does it crash? ::: No — range(0) is empty, input and output layers still run. nn.Linear(5,3) parameter count? ::: $5 \times 3 + 3 = 18$ (weight $3 \times 5$ + bias $3$ ). Dropout in eval mode? ::: Identity — deterministic pass-through, no zeroing. Two Linear layers sharing one weight tensor — counted how many times? ::: Once; .parameters() de-duplicates.