# nn.Module: the shape of every model — PyTorch

Source: https://www.skillbyai.com/en/pytorch/pt-nn-module

> Write an nn.Module subclass with __init__ and forward, count its parameters, and avoid the plain-list registration bug.

## 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.

![Diagram of 4 connected stages from left to right: input batch, layers (Linear, activation), logits, loss function.](assets/figures/pytorch/section-4-map.svg) — Figure 4.1: input batch -> layers (Linear, activation) -> logits -> loss function

## 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.

```python
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.

**Quiz:** Why should you call `model(x)` rather than `model.forward(x)`?

- [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

*Answer:* `model(x)` runs registered hooks around `forward`, which calling `forward` directly skips. `nn.Module.__call__` handles hooks and then calls `forward`. Calling `forward` directly bypasses that machinery; gradients and devices are unaffected.