nn.Module

nn.Module is the base class every PyTorch model, layer, and even loss function inherits from — the single most important building block for writing PyTorch code, used implicitly in nearly every code example throughout this hub.

The Standard Pattern

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()   # ALWAYS call this first -- registers the module properly
        self.layer1 = nn.Linear(10, 32)
        self.layer2 = nn.Linear(32, 1)

    def forward(self, x):
        x = torch.relu(self.layer1(x))
        return self.layer2(x)

model = MyModel()
output = model(x)   # calling model(x) automatically invokes forward(x) -- never call .forward() directly

What super().__init__() Actually Does

It initializes internal bookkeeping nn.Module needs to track every submodule and parameter you assign as an attribute — this is exactly what makes self.layer1 = nn.Linear(...) automatically get registered and discoverable via .parameters(), without you needing to manually track it anywhere.

Accessing Parameters

for name, param in model.named_parameters():
    print(name, param.shape)
# layer1.weight torch.Size([32, 10])
# layer1.bias torch.Size([32])
# layer2.weight torch.Size([1, 32])
# layer2.bias torch.Size([1])

total_params = sum(p.numel() for p in model.parameters())

train() and eval() Modes

model.train()   # activates dropout, batch norm uses BATCH statistics -- for training
model.eval()     # deactivates dropout, batch norm uses RUNNING statistics -- for validation/inference

This is exactly the mechanism flagged in Validation Loop — every nn.Module tracks a training/eval mode flag, and layers like nn.Dropout and nn.BatchNorm2d check this flag to change their behavior accordingly.

Common Mistakes

• Forgetting super().__init__() — this causes cryptic errors, since nn.Module's internal parameter-tracking machinery never gets initialized.
• Calling model.forward(x) directly instead of model(x) — calling the model instance directly triggers additional important internal machinery (hooks, mode-dependent behavior) that calling .forward() directly bypasses.

Interview Relevance

Q: "Why should you call model(x) rather than model.forward(x) directly in PyTorch?" Calling the model instance directly (model(x)) invokes nn.Module's __call__ method, which runs important additional machinery — like registered forward hooks — before and after actually calling forward(). Calling .forward() directly bypasses this machinery, which can silently break functionality that depends on it.

Practice Question

Why does assigning a layer as self.layer1 = nn.Linear(10, 32) inside __init__ automatically make its parameters show up in model.parameters(), with no extra code needed?

Related DL Notes

• PyTorch Layers
• PyTorch Training Loop
• PyTorch Custom Training Loops