PyTorch nn.Module Basics

nn.Module is the base class for every neural network in PyTorch. When you build any model — a simple classifier, a CNN, a transformer — you inherit from nn.Module. It provides the plumbing that makes models work: parameter tracking, GPU movement, saving and loading, and more.

Why nn.Module Exists

Without nn.Module, you would manage all model parameters manually — storing them in variables, tracking them for the optimizer, moving them to GPU one by one. nn.Module handles all of that automatically. Every layer you add inside an nn.Module subclass gets registered and tracked.

Building Your First nn.Module

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()           # always call this first
        self.layer1 = nn.Linear(4, 8)   # 4 inputs → 8 outputs
        self.layer2 = nn.Linear(8, 1)   # 8 inputs → 1 output

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

model = SimpleModel()
print(model)

Output:

SimpleModel(
  (layer1): Linear(in_features=4, out_features=8, bias=True)
  (layer2): Linear(in_features=8, out_features=1, bias=True)
)

Two Required Parts

__init__(self):
  ← define all layers here
  ← PyTorch registers them automatically
  ← must call super().__init__() first

forward(self, x):
  ← define how data flows through layers
  ← called when you do model(input)
  ← return the output

The forward Method

The forward method defines the data path. You write it like normal Python code. PyTorch calls it when you pass data to the model, and autograd records every operation for gradient computation.

class TwoLayerNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 20)
        self.fc2 = nn.Linear(20, 5)

    def forward(self, x):
        out = self.fc1(x)       # linear transformation
        out = torch.relu(out)   # activation
        out = self.fc2(out)     # second linear layer
        return out

Running Data Through the Model

model = TwoLayerNet()

# Sample batch: 8 samples, each with 10 features
X = torch.rand(8, 10)

output = model(X)   # internally calls model.forward(X)
print(output.shape)  # torch.Size([8, 5])

Data Flow Diagram

Input X: shape [8, 10]
        │
        ▼
   fc1: Linear(10→20)
        │
        ▼
   ReLU activation
        │
        ▼
   fc2: Linear(20→5)
        │
        ▼
Output: shape [8, 5]   ← 8 samples, each with 5 class scores

Accessing Model Parameters

model = SimpleModel()

# All parameters (weights + biases of all layers)
for name, param in model.named_parameters():
    print(f"{name}: shape={param.shape}, requires_grad={param.requires_grad}")

# Count total parameters
total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {total_params}")

Moving the Model to GPU

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleModel().to(device)

# All parameters and buffers automatically move to GPU
print(next(model.parameters()).device)   # cuda:0

train() and eval() Modes

Models behave differently during training vs inference. Batch normalization and dropout (both covered later) act differently depending on the mode. Always switch modes explicitly.

model.train()    # training mode: dropout and BN behave as in training
model.eval()     # evaluation mode: dropout disabled, BN uses running stats

Nested Modules

You can use one nn.Module inside another. This lets you build complex architectures from reusable pieces.

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(28*28, 128)

    def forward(self, x):
        return torch.relu(self.fc(x))

class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()      # nested module
        self.output = nn.Linear(128, 10)

    def forward(self, x):
        x = self.encoder(x)
        return self.output(x)

Summary

nn.Module is the foundation of every PyTorch model. Subclass it, define your layers in __init__, and describe the data flow in forward. PyTorch automatically tracks all parameters inside registered layers. Use model.parameters() to access weights for the optimizer. Call model.to(device) to move the model to GPU. Switch between model.train() and model.eval() to control training-specific behavior. Nest modules inside modules to build complex architectures cleanly.