nn.Module: the shape of every model

What nn.Module does for you

Every PyTorch model, from a two-layer network to a large language model, is a subclass of nn.Module. The pattern has two parts. In __init__, call super().__init__() and then create layers as attributes (self.fc1 = nn.Linear(...)). Assigning a module or an nn.Parameter to an attribute registers it, so its weights are found automatically. In forward(self, x), describe how input flows through those layers.

Call the model as model(x), not model.forward(x). The __call__ method runs any registered hooks and then forward, and some tools rely on those hooks.

Registration powers everything else: model.parameters() yields every learnable tensor for the optimizer, model.to(device) moves them all, model.state_dict() saves them, and model.train() or model.eval() reaches every sub-layer. The catch: layers stored in a plain Python list or dict are not registered. Use nn.ModuleList or nn.ModuleDict instead, or the optimizer will silently never update those layers.

From input to loss

An input batch flows through the layers defined in the module, producing raw output scores called logits. A loss function compares the logits with the targets to produce the number that training minimises.

A small multilayer perceptron

The parameter count is easy to verify by hand: 784 x 128 weights + 128 biases + 128 x 10 weights + 10 biases = 101770.

import torch
from torch import nn

class MLP(nn.Module):
    def __init__(self, in_features: int, hidden: int, num_classes: int):
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden)
        self.act = nn.ReLU()
        self.fc2 = nn.Linear(hidden, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(self.act(self.fc1(x)))

model = MLP(784, 128, 10)
x = torch.randn(32, 784)          # a batch of 32 flattened 28x28 images
logits = model(x)
print(logits.shape)               # torch.Size([32, 10])
print(sum(p.numel() for p in model.parameters()))   # 101770

The plain-list bug

self.layers = [nn.Linear(64, 64) for _ in range(4)] looks fine and even runs, but those layers are invisible to parameters(), to(device) and state_dict(). Write self.layers = nn.ModuleList([...]). If your parameter count looks too small, this is the first thing to check.

Quick check: Why should you call `model(x)` rather than `model.forward(x)`?

• `model(x)` runs registered hooks around `forward`, which calling `forward` directly skips
• `model.forward(x)` does not compute gradients
• `model(x)` automatically moves `x` to the GPU
• `model.forward(x)` always runs in evaluation mode