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.

Initialization on Programmer.ie

The Network: What Is It Without nn.Module?