nn.Module Explained: The Same Model Built with Raw Tensors and with nn.Module

Every PyTorch model is a subclass of nn.Module. This article explains what that class does by building the same model two ways: first with raw tensors only, then with nn.Module. Comparing the two shows exactly which parts of the work the module takes over.

The problem

The model learns the line y = 3x + 1 from 20 noisy points.

import torch

torch.manual_seed(0)
x = torch.linspace(0, 2, 20).unsqueeze(1)      # shape (20, 1)
y = 3 * x + 1 + 0.3 * torch.randn_like(x)      # shape (20, 1)

Version 1: raw tensors

A linear model has two parameters, a weight w and a bias b. Here they are created as plain tensors with requires_grad=True, so autograd tracks them.

w = torch.randn(1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
lr = 0.1

for epoch in range(200):
    y_hat = w * x + b                      # 1. forward pass
    loss = ((y_hat - y) ** 2).mean()       # 2. loss (MSE)
    loss.backward()                        # 3. gradients into w.grad, b.grad
    with torch.no_grad():                  # 4. update
        w -= lr * w.grad
        b -= lr * b.grad
    w.grad.zero_()                         # 5. reset gradients
    b.grad.zero_()

print(f"y = {w.item():.4f}x + {b.item():.4f}")

Output:

y = 3.2885x + 0.8159

The result is not exactly 3x + 1 because of the noise in the data. The best straight line through these 20 points, computed directly with least squares (torch.linalg.lstsq), is y = 3.2899x + 0.8142. Gradient descent reached that line to within 0.002.

The loop has five steps: forward pass, loss, backward pass, update, reset. Steps 4 and 5 each need a short explanation.

Why the update is inside torch.no_grad()

w is a leaf tensor with requires_grad=True. PyTorch does not allow a tracked leaf to be modified in place, because the computation graph depends on its value. Without no_grad(), the first update raises an error:

w -= lr * w.grad

RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.

Writing the update as w = w - lr * w.grad avoids that error but creates a different problem. The result of w - lr * w.grad is a new, non-leaf tensor, and autograd only stores .grad on leaf tensors. On the next step w.grad is None:

TypeError: unsupported operand type(s) for *: 'float' and 'NoneType'

Inside torch.no_grad(), the in-place update is allowed and is not recorded in the graph. w stays the same leaf tensor.

Why the gradients are reset

loss.backward() adds to .grad. It does not overwrite it. If the gradients are not reset, each step uses the sum of all previous gradients. With the same data and seed, the loss over time looks like this:

epoch                 0       2       5      10      50     199
----------------------------------------------------------------
with zero_()      31.47    3.70    0.66    0.42    0.09   0.067
without zero_()   31.47    1.13   24.47   16.28   29.91   33.73

Without the reset, the updates grow and change direction, and the loss does not converge.

Version 2: nn.Module

The same model, written as an nn.Module subclass:

import torch.nn as nn

class TinyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)

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

model = TinyModel()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.MSELoss()

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

w = model.linear.weight.item()
b = model.linear.bias.item()
print(f"y = {w:.4f}x + {b:.4f}")

Output:

y = 3.2882x + 0.8163

The result matches Version 1. The small difference in the last digits comes from a different random starting value.

Each part of Version 2 replaces a part of Version 1:

Version 1 (raw tensors)	Version 2 (nn.Module)
w, b created with requires_grad=True	nn.Linear(1, 1)
y_hat = w * x + b	model(x)
((y_hat - y) ** 2).mean()	nn.MSELoss()
update inside torch.no_grad()	optimizer.step()
w.grad.zero_(), b.grad.zero_()	optimizer.zero_grad()

nn.Linear does no extra computation. Its forward pass is x @ W.T + b:

lin = nn.Linear(1, 1)
xx = torch.randn(5, 1)
print(torch.allclose(lin(xx), xx @ lin.weight.T + lin.bias))    # True

Its weight and bias are ordinary leaf tensors with requires_grad=True, of type nn.Parameter:

for name, p in model.named_parameters():
    print(name, tuple(p.shape), p.is_leaf, type(p).__name__)

linear.weight (1, 1) True Parameter
linear.bias (1,) True Parameter

What nn.Module keeps track of

nn.Module keeps a record of the parameters and sub-modules that belong to the model. This record is filled automatically when attributes are assigned in __init__.

nn.Module overrides attribute assignment (__setattr__). When a line such as self.linear = nn.Linear(1, 1) runs, the module checks the type of the value:

• if it is an nn.Module, it is stored in the module's _modules dictionary
• if it is an nn.Parameter, it is stored in the _parameters dictionary
• anything else is stored as an ordinary Python attribute

print(nn.Linear(1, 1)._parameters.keys())    # odict_keys(['weight', 'bias'])

This is why super().__init__() must be the first line of __init__. It creates these dictionaries. If it is missing, the first assignment fails:

AttributeError: cannot assign module before Module.__init__() call

The main methods of nn.Module all work by going through these dictionaries, recursively into every sub-module:

Method	What it does with the registered parameters
model.parameters()	returns them, usually to pass to the optimizer
model.to(device)	moves each one to the device
model.state_dict()	collects them for saving
print(model)	lists the registered sub-modules

None of these methods read forward(). A layer is included only if it was registered in __init__.

Storing layers in a list: nn.ModuleList

A Python list is not an nn.Module, so a list assigned to self is stored as an ordinary attribute, and the layers inside it are not registered.

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.blocks = [nn.Linear(10, 64), nn.Linear(64, 64), nn.Linear(64, 64)]   # plain list
        self.head = nn.Linear(64, 3)

    def forward(self, x):
        for layer in self.blocks:
            x = torch.relu(layer(x))
        return self.head(x)

forward() still uses all four layers, so the model runs without error. But only head is registered:

print(sum(p.numel() for p in model.parameters()))    # 195
print(model)

Net(
  (head): Linear(in_features=64, out_features=3, bias=True)
)

The full network has (10*64 + 64) + 2*(64*64 + 64) + (64*3 + 3) = 9219 parameters. The optimizer receives 195 of them, so the three hidden layers are never updated. Their gradients are still computed during backward(), because they are part of the computation, but the optimizer does not apply them.

The fix is nn.ModuleList, which is itself an nn.Module and registers every layer inside it:

self.blocks = nn.ModuleList([nn.Linear(10, 64), nn.Linear(64, 64), nn.Linear(64, 64)])

Trained on the same 3-class dataset (2000 training rows, Adam, lr=1e-3), the two versions give:

                        parameters trained   test acc, epoch 10     50       300
--------------------------------------------------------------------------------
nn.ModuleList([...])        9219 / 9219              55.5%        84.9%    84.4%
plain list [...]             195 / 9219              36.2%        53.6%    61.3%

The unregistered layers are also missing from state_dict(), so they are not saved in a checkpoint:

nn.ModuleList: blocks.0.weight, blocks.0.bias, blocks.1.weight, blocks.1.bias,
               blocks.2.weight, blocks.2.bias, head.weight, head.bias
plain list:    head.weight, head.bias

And model.to(...) does not move them. After model.to(torch.float64), the head is float64 and the hidden layers remain float32, and the next forward pass fails:

RuntimeError: mat1 and mat2 must have the same dtype, but got Double and Float

The same rule applies to other containers and to raw tensors:

What you want to store	Use
a list of layers	nn.ModuleList
a dictionary of layers	nn.ModuleDict
layers applied one after another	nn.Sequential
a single trainable tensor	nn.Parameter(tensor)

A quick check for any model is to compare sum(p.numel() for p in model.parameters()) with the count calculated from the layer sizes.

Layers belong in __init__, not in forward

A layer created inside forward() is a new layer, with new random weights, on every call:

def forward(self, x):
    layer = nn.Linear(10, 64)          # new layer each call
    return self.head(torch.relu(layer(x)))

Passing the same input twice gives two different outputs:

[[ 0.326 -0.136 -0.508]]
[[-0.650 -0.677 -0.995]]

Layers are created once in __init__ and used in forward().

model(x) vs model.forward(x)

Calling model(x) runs nn.Module.__call__, which calls forward() and also runs any hooks registered on the module. Calling model.forward(x) runs only forward() and skips the hooks.

calls = []
model.register_forward_hook(lambda mod, inp, out: calls.append("hook"))

model(x);          print(calls)    # ['hook']
model.forward(x);  print(calls)    # ['hook']   (the hook did not run again)

Hooks are used by tools that attach to a model, such as pruning (torch.nn.utils.prune), quantization observers, feature extractors and FLOP counters. Always call model(x).

Summary

• A model is a set of parameters and a forward function. It can be written with raw tensors.
• With raw tensors, the update must be inside torch.no_grad() and the gradients must be reset each step.
• nn.Module registers every nn.Module and nn.Parameter assigned as an attribute in __init__. super().__init__() must come first.
• .parameters(), .to(), state_dict() and print() only see registered parameters and sub-modules.
• Layers in a plain Python list are not registered. Use nn.ModuleList, nn.ModuleDict or nn.Sequential.
• Create layers in __init__, use them in forward(), and call the model as model(x).

---

This is one chapter's worth of an idea from my book, PyTorch From Ground Up, which builds everything from tensors upward so nothing stays vague. If it helped: 8 chapters are free, no email required, there's a free one-page tensor cheat-sheet here, every example runs in the companion notebooks on GitHub, and the full book is on Leanpub or in paperback and Kindle on Amazon.

---

More in this series

How Training Actually Works:

• PyTorch Autograd Explained: What .backward() Actually Does
• Backpropagation by Hand: Two Layers, a Pen, and Then Autograd Agrees
• What optimizer.step() Actually Does: SGD, Momentum and Adam by Hand
• What Your Loss Function Actually Tells the Model

Shape mechanics:

• Reshape vs View in PyTorch
• PyTorch Broadcasting Explained
• What Does unsqueeze Do in PyTorch?
• What Does keepdim Do in PyTorch?
• PyTorch permute vs transpose

Next in this series: model.train(), model.eval() and torch.no_grad().