01 · The `nn.Module` Basics 🧱

4 min read

Every model you ever build in PyTorch — from a two-line regressor to a giant Transformer — is a subclass of nn.Module. Learn this pattern once, reuse it forever.

1. The 'Why'

In Module 02 you updated a single weight by hand: create a tensor with requires_grad=True, compute a loss, call .backward(), nudge the value. That's fine for one parameter, but a real network has thousands or millions of them, spread across many layers. Managing that by hand — keeping track of every weight tensor, remembering to move each one to the GPU, collecting them all to hand to an optimizer — would be a nightmare of bookkeeping. You'd spend all your time plumbing and none of it modeling.

nn.Module is PyTorch's answer to that bookkeeping problem. It's a base class you subclass to define your model, and it quietly handles the tedious parts: it discovers every parameter and sub-layer you register, lets you move the whole model to a device in one call, toggles training vs. evaluation behavior, and makes saving/loading trivial. You focus on two things — what layers your model contains (declared in __init__) and how data flows through them (declared in forward) — and nn.Module takes care of the rest. This separation of "structure" from "computation" is the single most important pattern in all of PyTorch.

2. Core Concepts

Subclassing nn.Module. You define a model as a class that inherits from nn.Module. Two methods matter: __init__ (declare your layers) and forward (define the computation).

Registering layers in __init__. Assigning a layer to self (e.g. self.fc = nn.Linear(...)) registers it. PyTorch then automatically tracks its parameters. Always call super().__init__() first, or registration silently breaks.

The forward method. This describes how an input tensor becomes an output — the actual math. You write it as ordinary Python; autograd builds the graph as it runs.

Call the module, not forward directly. You invoke a model as model(x), not model.forward(x). The __call__ machinery runs important hooks around your forward; calling forward directly skips them.

3. Code in Action

A minimal module

import torch
import torch.nn as nn

class LinearRegressor(nn.Module):
    def __init__(self):
        super().__init__()                 # MUST come first — sets up nn.Module
        self.linear = nn.Linear(1, 1)      # one input feature -> one output

    def forward(self, x):                  # defines the computation
        return self.linear(x)              # pass input through the layer

model = LinearRegressor()
x = torch.tensor([[2.0]])                  # a batch of 1 sample, 1 feature
y = model(x)                               # call the MODEL (not model.forward)
print(y)                                   # a 1x1 prediction tensor

Registration in action

import torch.nn as nn

model = LinearRegressor()
# Because self.linear was registered, its parameters are discoverable:
for name, param in model.named_parameters():
    print(name, tuple(param.shape))
# linear.weight (1, 1)
# linear.bias   (1,)

Why you call the module, not `forward`

import torch

x = torch.tensor([[1.0]])
y1 = model(x)          # ✅ correct: runs __call__ -> hooks -> forward
y2 = model.forward(x)  # ⚠️ works numerically but skips hooks; avoid this

4. Common Pitfalls

Pitfall 1 — Forgetting super().__init__(). If you omit it (or place it after your layer assignments), nn.Module isn't initialized and your layers won't register — model.parameters() comes back empty and the optimizer has nothing to train. Always make it the first line of __init__.

Pitfall 2 — Calling model.forward(x) directly. It usually returns the right numbers, so the bug hides — but it bypasses hooks that some features (and later modules like quantization) depend on. Always call model(x).

Pitfall 3 — Storing layers in a plain Python list. self.layers = [nn.Linear(...), ...] does not register the layers, so their parameters are invisible to the optimizer. Use nn.ModuleList or nn.Sequential instead (covered in sub-module 03).

5. Further Reading & Watch List

➡️ Next: 02 · Common Layers