Building a Neural Network

Build the mental model

PyTorch's `nn.Module` is the base class every network builds on: you declare its layers as attributes inside `__init__` and describe how data flows through them inside `forward()`. `nn.Linear(in_features, out_features)` is the workhorse layer — a fully-connected layer computing `x @ W^T + b`, where W and b are learnable parameters PyTorch initializes automatically. You could, in principle, write this same computation as raw tensor operations with manually created weight tensors, but subclassing `nn.Module` buys you several things for free that raw tensor math doesn't: PyTorch automatically discovers every parameter you declare (so `model.parameters()` hands the optimizer everything that needs updating, without you enumerating tensors by hand), `model.to(device)` moves every parameter to a GPU in one call, `torch.save`/`load_state_dict` handles saving and restoring all weights, and modules compose — a bigger network can simply contain smaller `nn.Module`s as sub-layers. This composability is why real architectures, from a two-layer classifier to a multi-block transformer, are all built from the same basic pattern.

Connect it to a real scenario

The sentiment classifier the Tutorial Platform eventually ships — taking a feedback comment's numeric representation and predicting positive/negative — is exactly an nn.Module: an input layer sized to match the comment's feature vector, one or more hidden nn.Linear layers, and an output layer producing a single score. Defining it this way means the same class can later be moved to a GPU, saved after training, and swapped for a bigger architecture (a CNN or transformer in later chapters) without changing how the rest of the platform's training or serving code calls it.

Declare layers in __init__ and the data flow in forward() when subclassing nn.Module.
nn.Linear(in, out) is a fully-connected layer computing x @ W^T + b with learnable W and b.
Subclassing nn.Module gives automatic parameter tracking, .to(device), save/load, and composability for free.

Try the working example

python

import torch
import torch.nn as nn

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(10, 5)  # 10 input features -> 5 hidden units
        self.layer2 = nn.Linear(5, 1)   # 5 hidden units -> 1 output

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

model = TinyNet()
sample_input = torch.randn(4, 10)  # batch of 4 examples, 10 features each
output = model(sample_input)

print("Output shape:", output.shape)
print("Number of parameters:", sum(p.numel() for p in model.parameters()))

You should see

Prints Output shape: torch.Size([4, 1]) — one prediction per example in the batch of 4 — and Number of parameters: 61 (55 from layer1's weights+bias, 6 from layer2's).

5-minute try-it

Add a third nn.Linear layer to TinyNet (e.g. 5 -> 5 -> 1 becomes 5 -> 3 -> 1), re-run the forward pass, and print the new total parameter count to see how it changes.

One important caution

Forgetting to call super().__init__() before assigning layers — nn.Module relies on its own __init__ to set up internal bookkeeping, and skipping it means your layers silently won't show up in model.parameters().

Confusing the input feature dimension between layers (e.g. layer2 expecting an input size that doesn't match layer1's output size) raises a shape-mismatch RuntimeError at the first forward pass, not at model definition time.

PyTorch Docs — Build the Neural Network — Deep Learning