4. Model architecture definition with nn.Module

In PyTorch, every non-trivial model is built by subclassing torch.nn.Module. This base class is far more than a naming convention: it is the mechanism that registers parameters, organizes submodules, and defines how tensors flow through the network during a forward pass. Understanding what nn.Module actually does under the hood is the difference between writing models that train reliably and models that silently fail to update some of their weights.

What nn.Module Really Manages

When you assign a nn.Parameter or another nn.Module as an attribute of a module, PyTorch intercepts that assignment through __setattr__ and records the object in internal dictionaries: _parameters for learnable tensors and _modules for nested modules. This registration is what makes model.parameters(), model.named_parameters(), model.state_dict(), and model.to(device) work recursively. A plain tensor assigned as an attribute is not registered — it will not appear in state_dict(), will not move to GPU with .to(device), and will not receive gradients. This is one of the most common sources of subtle bugs.

Layers, Containers, and Parameter Registration

PyTorch provides two broad categories of building blocks. Layers such as nn.Linear, nn.Conv2d, and nn.LayerNorm hold their own parameters and implement a forward method. Containers such as nn.Sequential, nn.ModuleList, and nn.ModuleDict hold other modules and manage their registration. The distinction matters: nn.Sequential also defines a forward pass (chaining its children in order), whereas nn.ModuleList and nn.ModuleDict are pure containers with no forward logic — you must iterate over them yourself.

A frequent mistake is using a Python list to hold submodules. A plain list is not registered, so its contents never appear in parameters(). The correct pattern is:

class MLP(nn.Module):
    def __init__(self, in_dim, hidden, out_dim, depth=3):
        super().__init__()
        dims = [in_dim] + [hidden] * depth + [out_dim]
        self.layers = nn.ModuleList(
            nn.Linear(dims[i], dims[i + 1]) for i in range(len(dims) - 1)
        )
        self.act = nn.GELU()

    def forward(self, x):
        for layer in self.layers[:-1]:
            x = self.act(layer(x))
        return self.layers[-1](x)

Here nn.ModuleList ensures every nn.Linear is registered, so model.parameters() yields all weight and bias tensors. If you replaced it with a Python list, the optimizer would receive an empty parameter set and training would appear to run while the loss never decreases.

Defining the Forward Computation Graph

The forward method defines the computation graph. PyTorch builds this graph dynamically at each call using autograd, recording every operation on tensors that require gradients. The graph is not stored on the module; it is constructed per forward pass and discarded after backward(). This is why control flow — loops, conditionals, even recursion — is allowed inside forward, unlike static-graph frameworks.

Two rules keep the graph correct. First, never call forward directly; always call model(x), which invokes __call__ and runs all registered hooks. Second, keep tensor operations inside forward; doing heavy computation in __init__ on parameters is fine, but any tensor that should participate in autograd must be produced during the forward pass.

Worked Example: A Small Residual Block

Consider a residual block for a 1D signal with 64 channels. The block applies two convolutions with normalization and adds the input to the output. The skip connection requires the input and output shapes to match; if they differ, a projection is needed.

class ResidualBlock(nn.Module):
    def __init__(self, channels=64, kernel_size=3):
        super().__init__()
        padding = kernel_size // 2
        self.conv1 = nn.Conv1d(channels, channels, kernel_size, padding=padding)
        self.norm1 = nn.BatchNorm1d(channels)
        self.conv2 = nn.Conv1d(channels, channels, kernel_size, padding=padding)
        self.norm2 = nn.BatchNorm1d(channels)
        self.act = nn.ReLU(inplace=True)

    def forward(self, x):
        identity = x
        out = self.act(self.norm1(self.conv1(x)))
        out = self.norm2(self.conv2(out))
        return self.act(out + identity)

With channels=64 and kernel_size=3, the padding of 1 preserves the temporal length, so out + identity is shape-compatible. A quick sanity check confirms registration and shape:

block = ResidualBlock()
x = torch.randn(8, 64, 128)   # batch=8, channels=64, length=128
y = block(x)
print(y.shape)                 # torch.Size([8, 64, 128])
print(len(list(block.parameters())))  # 8 parameter tensors

The eight parameters come from two convolutions (weight + bias each) and two batch-norm layers (weight + bias each). If the count is lower, a layer was not registered. If the forward call raises a shape error, the padding or channel count is inconsistent.

Visualizing the Registration and Forward Flow

The diagram below shows how a module's attributes map to registered parameters and submodules, and how the forward pass traverses them to build the autograd graph.

Common Pitfalls and How to Detect Them

Three failure modes recur in practice. First, unregistered parameters: a tensor assigned directly as an attribute never reaches the optimizer. Symptom: loss plateaus immediately. Detection: compare len(list(model.parameters())) against the expected count. Second, shared modules: assigning the same module instance to two attributes registers it once, so gradients accumulate from both paths — usually intended, but surprising if not. Third, device mismatch: buffers created on CPU stay on CPU unless registered, causing runtime errors when the model is moved to CUDA.

A reliable verification routine after defining any module is to instantiate it, run a dummy forward pass with the expected input shape, and inspect state_dict().keys(). If the keys match the layers you defined and the output shape matches your expectation, registration and the forward graph are consistent. This check takes seconds and catches the majority of architecture-definition bugs before training begins.