What is nn.Module?

nn.Module is the base class for all neural network models.

What is nn.Module?

nn.Module is the base class every model and layer in PyTorch derives from. A module keeps its parameters, its submodules and the forward computation together:

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.katman = nn.Linear(4, 3)

    def forward(self, x):
        return self.katman(x)

m = Model()
print(m(torch.zeros(2, 4)).shape)   # torch.Size([2, 3])

A module is called as m(x); behind the scenes this runs forward.

Ready-made modules

Inside torch.nn, layers (Linear, Conv2d, LSTM), activations (ReLU, Sigmoid), losses (MSELoss, CrossEntropyLoss) and regularisers (Dropout, BatchNorm1d) are all modules. For simple chains, nn.Sequential builds a model without writing a class:

model = nn.Sequential(
    nn.Linear(4, 16),
    nn.ReLU(),
    nn.Linear(16, 3),
)
print(model(torch.zeros(5, 4)).shape)   # torch.Size([5, 3])

Exercise First module

Build the model nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3)). Print the output shape for the input torch.zeros(5, 4) and the model's length (len(model)) on one line.

Show example solution

import torch
from torch import nn

model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))
print(model(torch.zeros(5, 4)).shape, len(model))

Test yourself

Why is <code>m(x)</code> written instead of <code>m.forward(x)</code> to run the model?
- m(x) also runs hooks and internal machinery
- Because it is shorter
- There is no forward
- There is no difference
Show answer

m(x) also runs hooks and internal machinery — __call__ wraps forward; calling forward directly skips these steps.

Chain the layers in order.

model = nn.____(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))

Show answer

model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))

Match the modules.

nn.Linear

Linear layer
nn.ReLU

Activation
nn.CrossEntropyLoss

Loss
nn.Sequential

Sequential container

Next lesson: __init__

We Are Enhancing Your Experience

Cookie Preferences

Strictly Necessary Cookies

Analytical Cookies