Building Models with nn.Module

nn.Module is the base class for all PyTorch models. Subclass it to define layers, implement forward(), and register parameters that optimizers update during training.

Recipe

Quick-reference recipe card - copy-paste ready.

import torch.nn as nn
 
class MLP(nn.Module):
    def __init__(self, in_dim: int, hidden: int, out_dim: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, out_dim),
        )
 
    def forward(self, x):
        return self.net(x)

When to reach for this:

• Defining any neural network architecture in PyTorch.
• Composing layers with custom forward logic (skip connections, attention).
• Inspecting, freezing, or selectively training parameter groups.
• Serializing models via state_dict().

Working Example

"""building_models.py - custom CNN for image classification."""
from __future__ import annotations
 
import torch
import torch.nn as nn
import torch.nn.functional as F
 
class SmallCNN(nn.Module):
    def __init__(self, num_classes: int = 10):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, num_classes)
        self.dropout = nn.Dropout(0.25)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.pool(F.relu(self.conv1(x)))   # 28->14
        x = self.pool(F.relu(self.conv2(x)))   # 14->7
        x = x.flatten(1)
        x = self.dropout(F.relu(self.fc1(x)))
        return self.fc2(x)
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SmallCNN(num_classes=10).to(device)
 
x = torch.randn(8, 1, 28, 28, device=device)
logits = model(x)
print("logits shape:", logits.shape)  # (8, 10)
 
total_params = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"params: {total_params:,} trainable: {trainable:,}")

What this demonstrates:

• Layers registered in __init__ become part of model.parameters().
• forward defines the computation; call via model(x), not model.forward(x).
• .to(device) moves all parameters and buffers to GPU.
• flatten(1) preserves batch dimension while collapsing spatial dims.

Deep Dive

How It Works

• __init__ registers submodules (nn.Linear, nn.Conv2d) and parameters (nn.Parameter).
• forward runs the computation graph; hooks can intercept inputs/outputs.
• nn.Module tracks training vs eval mode for dropout and batch norm behavior.
• state_dict() returns a dict of parameter tensors for save/load.
• named_modules() and named_children() traverse the module tree.

Module Patterns

Python Notes

# Freeze backbone, train head only
for param in model.conv1.parameters():
    param.requires_grad = False
 
# Different learning rates per group
optimizer = torch.optim.Adam([
    {"params": model.conv1.parameters(), "lr": 1e-5},
    {"params": model.fc2.parameters(), "lr": 1e-3},
])

Gotchas

• Calling forward() directly - skips hooks and nn.Module wrappers. Fix: always call model(x).
• Forgetting super().__init__() - submodules not registered. Fix: call super().__init__() first in __init__.
• Not calling model.eval() for inference - dropout and batch norm behave incorrectly. Fix: model.eval() before inference; model.train() for training.
• Creating layers in forward - new parameters every pass, never trained. Fix: define all layers in __init__.
• Device mismatch - model on GPU, input on CPU. Fix: model.to(device) and x.to(device).
• BatchNorm with batch_size=1 - statistics are undefined. Fix: use model.eval() or nn.GroupNorm for small batches.

Alternatives

FAQs

sum(p.numel() for p in model.parameters() if p.requires_grad)

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
 
model.apply(init_weights)

print(model)
# or
from torchinfo import summary
summary(model, input_size=(8, 1, 28, 28))

def forward(self, x):
    residual = x
    out = self.block(x)
    return F.relu(out + residual)

Related

• PyTorch Basics - tensor and nn fundamentals
• Training Loops - training the model
• Transfer Learning & Fine-Tuning - freezing layers
• Saving, Loading & Export - state_dict persistence
• PyTorch Lightning - LightningModule wrapper

Stack versions: This page was written for Python 3.14.0 (stable 3.14, maintenance 3.13), FastAPI 0.115+, Django 5.2, Flask 3.1, Pydantic 2, PyTorch 2.6+, pandas 2.2+, Polars 1.x, ruff 0.9+, and uv 0.6+. 3730363230327c3236366f6967637c6f692e73656469756765646f63