Building Models with nn.Module

Models are Python classes#

In PyTorch, a neural network is a class that inherits from nn.Module. You define layers in __init__ and describe how data flows through them in forward. PyTorch then tracks every parameter automatically.

Your first model#

Let's rebuild the XOR network from the Deep Learning tutorial — the one we wrote with NumPy — in PyTorch:

import torch
from torch import nn


class XORNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(2, 4)    # 2 inputs -> 4 hidden units
        self.output = nn.Linear(4, 1)    # 4 hidden units -> 1 output

    def forward(self, x):
        x = torch.tanh(self.hidden(x))
        return self.output(x)            # raw score (logit); no sigmoid here


model = XORNet()
print(model)
print("parameters:", sum(p.numel() for p in model.parameters()))

XORNet(
  (hidden): Linear(in_features=2, out_features=4, bias=True)
  (output): Linear(in_features=4, out_features=1, bias=True)
)
parameters: 17

17 parameters: (2 × 4 weights + 4 biases) + (4 × 1 weights + 1 bias).

Training it#

torch.manual_seed(0)
model = XORNet()

X = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=torch.float32)
y = torch.tensor([[0], [1], [1], [0]], dtype=torch.float32)

loss_fn = nn.BCEWithLogitsLoss()                         # sigmoid + binary cross-entropy, numerically stable
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

for epoch in range(500):
    optimizer.zero_grad()
    loss = loss_fn(model(X), y)
    loss.backward()
    optimizer.step()

predictions = (torch.sigmoid(model(X)) > 0.5).int().flatten().tolist()
print("predictions:", predictions)

predictions: [0, 1, 1, 0]

The whole manual backward pass from the NumPy version is now loss.backward(), and the four update lines are optimizer.step().

Common layers#

Activations: nn.ReLU(), nn.GELU(), nn.Tanh(), nn.Sigmoid(), nn.Softmax(dim=...).

nn.Sequential for simple stacks#

When data flows straight through, skip the class:

model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(128, 10),     # 10 class scores (logits)
)

Write a full nn.Module class when you need branching, skip connections, multiple inputs or custom logic in forward.

Choosing the loss function#

Don't add Softmax or Sigmoid before these "with logits" losses — they already include it, and doing it twice hurts training.

Inspecting a model#

for name, param in model.named_parameters():
    print(name, tuple(param.shape))

Helpful for debugging shapes, freezing specific layers, or checking what a fine-tuning method actually trains.

train() and eval() modes#

model.train() and model.eval() switch the behaviour of layers like Dropout and BatchNorm. Always call model.eval() before validation or inference, and model.train() before training again.

Common mistakes#

• Forgetting super().__init__() — PyTorch can't register your layers.
• Calling model.forward(x) directly — call model(x) so PyTorch's hooks run.
• Applying softmax before CrossEntropyLoss.
• Leaving the model in train mode for evaluation, so dropout makes results noisy.

Key takeaways#

• Define layers in __init__, data flow in forward, and call the model like a function.
• nn.Sequential is perfect for simple stacks; subclass nn.Module for anything custom.
• Match the final layer and loss: logits + BCEWithLogitsLoss or CrossEntropyLoss.
• Switch between train() and eval() modes correctly.