nn.Module — The Subclass Everything Inherits From

Two methods, infinite power

Every neural network component in PyTorch — from a single linear layer to a 70B-parameter transformer — is an nn.Module. Subclassing it is the single most important Python pattern you'll use in this framework, so internalize the contract:

Override __init__. Always call super().__init__() first. Create child modules and parameters by assigning them to self (e.g. self.layer1 = nn.Linear(10, 20)). PyTorch hooks the __setattr__ to register them.
Override forward(self, x, ...). Define the computation. Never call this directly — call the module instance like a function: model(x), not model.forward(x). The instance call routes through __call__, which runs registered hooks, handles training/eval modes for children, and threads through the autograd machinery properly.

What you get for free

model.parameters() — iterator over every learnable tensor, no matter how nested.
model.named_parameters() — same, with the dotted-path name attached.
model.to(device) — moves all parameters AND buffers to a device.
model.train() / model.eval() — flips behavior of children that care (Dropout, BatchNorm).
model.state_dict() / load_state_dict() — serialization-ready dict of all parameters and buffers.

All of this depends on you using nn.Module's machinery (assigning to self, using nn.ModuleList instead of plain Python lists for collections, using nn.Parameter for raw learnable tensors). Skip those, and the auto-magic stops working silently.

Code

Minimal nn.Module — the canonical subclass·python

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self, in_dim=784, hidden=128, out_dim=10):
        super().__init__()                    # ALWAYS first
        self.fc1 = nn.Linear(in_dim, hidden)
        self.act = nn.ReLU()
        self.fc2 = nn.Linear(hidden, out_dim)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        return x

model = TinyMLP()
print(model)
# TinyMLP(
#   (fc1): Linear(in_features=784, out_features=128, bias=True)
#   (act): ReLU()
#   (fc2): Linear(in_features=128, out_features=10, bias=True)
# )

The free utilities — parameters, devices, modes·python

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self): super().__init__(); self.fc = nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP()

# Parameter iteration
for name, p in model.named_parameters():
    print(name, p.shape)
# fc.weight torch.Size([4, 10])
# fc.bias   torch.Size([4])

# Total parameter count
total = sum(p.numel() for p in model.parameters())
print(f"Params: {total:,}")  # Params: 44

# Move once, everything follows
model = model.to('cpu')        # or 'cuda' / 'mps'
print(next(model.parameters()).device)

# Mode switching
model.train()                  # default mode
model.eval()                   # affects Dropout / BatchNorm

Why model(x), never model.forward(x)·python

import torch
import torch.nn as nn

class HookedLinear(nn.Linear):
    pass

m = HookedLinear(4, 2)

# Register a hook that prints the OUTPUT shape after each forward
def hook(module, inputs, output):
    print(f"Hook fired: out shape = {output.shape}")

m.register_forward_hook(hook)

x = torch.randn(3, 4)

m(x)               # Hook fired: out shape = torch.Size([3, 2])
# m.forward(x)     # NO hook fires! Bypasses __call__.

# This is why everyone calls model(x). Bypassing __call__ skips:
#   - registered forward / backward hooks
#   - __torch_function__ dispatch
#   - some compile / quantization machinery

nn.Module — The Subclass Everything Inherits From

Two methods, infinite power

What you get for free

Code

External links

Exercise

Progress

Comments 0