Nn.Module on AIBussin — AI applications, systems and books

The Network: What Is It Without nn.Module?

Sat, 08 Aug 2026 12:59:00 +0100

Here is a two-layer neural network, written with nothing but the machinery of the last three chapters. Four trainable tensors, a forward pass, a loss, backward(), an update. It runs without warnings, and the loss falls by a third.

import math
import torch

torch.manual_seed(0)

n = 2000
X = torch.randn(n, 2)
y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).long()

perm = torch.randperm(n)
n_train = int(0.8 * n)
X_train, y_train = X[perm[:n_train]], y[perm[:n_train]]
X_val, y_val = X[perm[n_train:]], y[perm[n_train:]]

D_in, H, C = 2, 16, 2

torch.manual_seed(1)
W1 = (torch.randn(D_in, H) / math.sqrt(D_in)).requires_grad_()
b1 = torch.zeros(H, requires_grad=True)
W2 = (torch.randn(H, C) / math.sqrt(H)).requires_grad_()
b2 = torch.zeros(C, requires_grad=True)
parameters = [W1, b1, W2, b2]

def forward(x):
 z1 = x @ W1 + b1
 return z1 @ W2 + b2

def cross_entropy(logits, targets):
 log_probs = torch.log_softmax(logits, dim=1)
 return -log_probs[torch.arange(targets.shape[0]), targets].mean()

def accuracy(logits, targets):
 return (logits.argmax(dim=1) == targets).float().mean().item()

learning_rate = 0.5

# First verify that one training step is mechanically healthy.
names = ["W1", "b1", "W2", "b2"]
initial_values = [p.detach().clone() for p in parameters]

loss = cross_entropy(forward(X_train), y_train)
before = [p.detach().clone() for p in parameters]

loss.backward()

with torch.no_grad():
 for p in parameters:
 p -= learning_rate * p.grad

for name, p, old in zip(names, parameters, before):
 print(f"{name:3s} leaf={p.is_leaf} grad_none={p.grad is None} "
 f"grad_norm={p.grad.norm().item():8.4f} "
 f"finite={torch.isfinite(p.grad).all().item()} "
 f"moved={(p.detach() - old).abs().max().item():.4f}")

W1 leaf=True grad_none=False grad_norm= 0.4516 finite=True moved=0.1048
b1 leaf=True grad_none=False grad_norm= 0.0192 finite=True moved=0.0053
W2 leaf=True grad_none=False grad_norm= 1.1058 finite=True moved=0.1918
b2 leaf=True grad_none=False grad_norm= 0.0175 finite=True moved=0.0062

Every parameter is a leaf. Every parameter requires gradients. Every parameter receives a finite gradient, and every parameter changes by a measurable amount when the update runs. The basic training machinery is working.

nn.Module: What Does PyTorch Think Belongs to Your Model?

Sat, 08 Aug 2026 13:09:00 +0100

Here is a classifier for the XOR-shaped dataset from Chapter 4, written the way most people write PyTorch. It trains. Validation accuracy climbs from chance to 95.7%. It saves a checkpoint, the checkpoint loads without a warning, and then the reloaded model performs at chance.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

n = 2000
X = torch.randn(n, 2)
y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).long()

perm = torch.randperm(n)
X_train, y_train = X[perm[:1600]], y[perm[:1600]]
X_val, y_val = X[perm[1600:]], y[perm[1600:]]

class Encoder(nn.Module):
 def __init__(self, width=16, depth=3):
 super().__init__()
 self.blocks = [nn.Linear(width, width) for _ in range(depth)]

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

class Classifier(nn.Module):
 def __init__(self, width=16, depth=3):
 super().__init__()
 self.stem = nn.Linear(2, width)
 self.encoder = Encoder(width, depth)
 self.head = nn.Linear(width, 2)

 def forward(self, x):
 x = torch.relu(self.stem(x))
 x = self.encoder(x)
 return self.head(x)

model = Classifier()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

for epoch in range(200):
 logits = model(X_train)
 loss = F.cross_entropy(logits, y_train)

 optimizer.zero_grad()
 loss.backward()
 optimizer.step()

 if epoch % 50 == 0 or epoch == 199:
 with torch.no_grad():
 val_acc = (model(X_val).argmax(1) == y_val).float().mean()
 print(f"epoch={epoch:3d} loss={loss.item():.4f} "
 f"val_acc={val_acc.item():.3f}")

torch.save(model.state_dict(), "classifier.pt")

reloaded = Classifier()
print(reloaded.load_state_dict(torch.load("classifier.pt", weights_only=True)))

with torch.no_grad():
 acc = (reloaded(X_val).argmax(1) == y_val).float().mean()
print(f"reloaded val_acc={acc.item():.3f}")

epoch= 0 loss=0.6964 val_acc=0.512
epoch= 50 loss=0.5930 val_acc=0.810
epoch=100 loss=0.3856 val_acc=0.873
epoch=150 loss=0.2550 val_acc=0.925
epoch=199 loss=0.1908 val_acc=0.957
<All keys matched successfully>
reloaded val_acc=0.512

Read the last two lines together. PyTorch reports that every key in the checkpoint matched every key the model expected, and the model that came back is worthless. Nothing raised. Nothing warned. load_state_dict was as satisfied as it is capable of being.

PyTorch nn.Module Explained: Missing Parameters, state_dict, Buffers and Registration Bugs

Sat, 08 Aug 2026 13:09:00 +0100

PyTorch: Zero to Hero — Step 04

In the previous post we built a neural network using raw tensors and autograd.

Now we are going to add the abstraction PyTorch expects almost every real model to use:

class Model(torch.nn.Module):
 ...

But this is not going to be a tour of nn.Module methods.

The useful question for programmers is:

What exactly does nn.Module register, track, move, save and expose — and how do those mechanisms break?

Build a Neural Network From Scratch in PyTorch Without nn.Module

Sat, 08 Aug 2026 12:59:00 +0100

PyTorch: Zero to Hero — Step 03

Most PyTorch tutorials begin with something like this:

import torch.nn as nn

model = nn.Sequential(
 nn.Linear(2, 16),
 nn.ReLU(),
 nn.Linear(16, 2),
)

That is useful code.

It is also hiding almost everything interesting.

In this post we are going to build the same kind of neural network using ordinary PyTorch tensors.

No nn.Module.

No nn.Linear.

No torch.optim.Adam.

No optimizer.step().

We will manually create the parameters, write the forward pass, calculate the loss, call autograd, update the weights, zero the gradients, batch the data, evaluate the model and then compare the result with the idiomatic PyTorch version.