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

AI, PyTorch
Page content

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.

The goal is not to avoid PyTorch abstractions forever.

The goal is to know what they are doing when we start using them.

What we are building

We will train a tiny classifier on a two-dimensional dataset.

Each input has two features:

x = [x1, x2]

The network will be:

    graph TD
    A["x [B,2]"] --> B["Linear(2→16)<br/>W1,b1"]
    B --> C["ReLU"]
    C --> D["Linear(16→2)<br/>W2,b2"]
    D --> E["logits [B,2]"]

Written mathematically:

hidden = ReLU(X @ W1 + b1)
logits = hidden @ W2 + b2

The trainable parameters are:

W1: [2, 16]
b1: [16]
W2: [16, 2]
b2: [2]

That is the entire network.

Everything else is training machinery.

1. Create a dataset

We want something nonlinear enough that a single linear classifier is not the whole story.

A simple XOR-like dataset works well.

import torch


torch.manual_seed(42)

n = 2000
X = torch.randn(n, 2)

y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).long()

print(X.shape)
print(y.shape)
print(X[:5])
print(y[:5])

Expected shapes:

torch.Size([2000, 2])
torch.Size([2000])

The label rule is:

class 1 when x1 and x2 have opposite signs
class 0 otherwise

A few points:

examples = torch.tensor([
    [ 1.0,  1.0],
    [ 1.0, -1.0],
    [-1.0,  1.0],
    [-1.0, -1.0],
])

example_y = ((examples[:, 0] > 0) ^ (examples[:, 1] > 0)).long()

print(example_y)

Output:

tensor([0, 1, 1, 0])

2. Split train and validation data

Do not evaluate your training code using only the data it has already seen.

indices = torch.randperm(n)

train_size = int(n * 0.8)
train_idx = indices[:train_size]
val_idx = indices[train_size:]

X_train = X[train_idx]
y_train = y[train_idx]

X_val = X[val_idx]
y_val = y[val_idx]

print(X_train.shape, y_train.shape)
print(X_val.shape, y_val.shape)

Output:

torch.Size([1600, 2]) torch.Size([1600])
torch.Size([400, 2]) torch.Size([400])

3. Create the weights manually

A neural network layer needs weights and usually a bias.

Let’s create them directly.

input_size = 2
hidden_size = 16
num_classes = 2

W1 = torch.randn(input_size, hidden_size) * 0.1
b1 = torch.zeros(hidden_size)

W2 = torch.randn(hidden_size, num_classes) * 0.1
b2 = torch.zeros(num_classes)

W1.requires_grad_()
b1.requires_grad_()
W2.requires_grad_()
b2.requires_grad_()

Inspect them:

for name, tensor in {
    "W1": W1,
    "b1": b1,
    "W2": W2,
    "b2": b2,
}.items():
    print(
        name,
        "shape=", tuple(tensor.shape),
        "requires_grad=", tensor.requires_grad,
        "is_leaf=", tensor.is_leaf,
    )

You want:

W1 shape=(2, 16)
b1 shape=(16,)
W2 shape=(16, 2)
b2 shape=(2,)

All four are leaf tensors that require gradients.

Those four tensors are our model.

There is no separate hidden object called a neural network.

A common initialization bug

This looks innocent:

W1 = torch.randn(2, 16, requires_grad=True) * 0.1

But now inspect it:

print(W1.is_leaf)

Depending on how you create and transform a tensor, the multiplication can make the resulting tensor non-leaf.

For manual parameter management, it is clearer to initialize first and then call:

W1.requires_grad_()

Or detach after initialization if necessary:

W1 = (torch.randn(2, 16) * 0.1).requires_grad_()

When you later use nn.Parameter, PyTorch gives parameters the semantics expected by nn.Module.

For now, we are doing it ourselves.

4. Write the first linear layer

The fundamental operation is matrix multiplication.

z1 = X_train @ W1 + b1

print(z1.shape)

Shape reasoning:

X_train: [1600, 2]
W1:      [2, 16]

[1600, 2] @ [2, 16]
             ↓
          [1600, 16]

The bias is:

[16]

PyTorch broadcasts it across the batch dimension.

So:

X_train @ W1 + b1

produces:

[1600, 16]

This is exactly the kind of shape reasoning from Step 01 that becomes unavoidable in neural-network code.

5. Add a nonlinear activation

If we stack only linear transformations, the result is still equivalent to one linear transformation.

We need a nonlinearity.

Let’s implement ReLU ourselves:

def relu(x):
    return torch.clamp_min(x, 0.0)

Now:

h = relu(X_train @ W1 + b1)

print(h.shape)

Output:

torch.Size([1600, 16])

You can also write:

h = torch.maximum(
    X_train @ W1 + b1,
    torch.tensor(0.0),
)

or:

h = (X_train @ W1 + b1).relu()

We will keep our wrapper because it makes the architecture obvious.

6. Write the output layer

The hidden representation goes into the second linear layer:

logits = h @ W2 + b2

print(logits.shape)

Output:

torch.Size([1600, 2])

Each row contains two scores:

[class_0_score, class_1_score]

These are called logits.

Do not apply softmax yet.

For classification training, PyTorch’s cross-entropy loss expects unnormalized logits.

7. Put the forward pass in a function

We now have enough code to define the model computation.

def forward(x, W1, b1, W2, b2):
    hidden = relu(x @ W1 + b1)
    logits = hidden @ W2 + b2
    return logits

Try it:

logits = forward(X_train, W1, b1, W2, b2)
print(logits.shape)

Output:

torch.Size([1600, 2])

This function is conceptually the equivalent of an nn.Module.forward() method.

There is nothing special about the forward pass itself.

It is ordinary tensor code.

8. Implement cross entropy manually

We could immediately use:

import torch.nn.functional as F
loss = F.cross_entropy(logits, y_train)

But once, it is worth seeing what classification loss is doing.

For numerical stability, we can use log_softmax:

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

Now:

logits = forward(X_train, W1, b1, W2, b2)
loss = cross_entropy(logits, y_train)

print(loss)

What happened here?

log_probs = torch.log_softmax(logits, dim=1)

turns each row of logits into log probabilities.

Then:

log_probs[torch.arange(targets.shape[0]), targets]

selects the log probability assigned to the correct class for each sample.

Finally:

.mean()

averages the negative log likelihood across the batch.

That is enough for our classifier.

Verify against PyTorch

We should not trust our implementation merely because it runs.

import torch.nn.functional as F

manual_loss = cross_entropy(logits, y_train)
torch_loss = F.cross_entropy(logits, y_train)

print("manual:", manual_loss.item())
print("torch: ", torch_loss.item())
print("close: ", torch.allclose(manual_loss, torch_loss))

The final line should be:

close: True

This pattern is valuable beyond tutorials:

Implement the thing you want to understand, then compare it against a trusted implementation.

9. Backpropagate

Now the interesting part becomes tiny:

loss.backward()

Inspect the gradients:

for name, tensor in {
    "W1": W1,
    "b1": b1,
    "W2": W2,
    "b2": b2,
}.items():
    print(
        name,
        "grad shape=", None if tensor.grad is None else tuple(tensor.grad.shape),
        "grad norm=", None if tensor.grad is None else tensor.grad.norm().item(),
    )

You should see gradients for every parameter.

The gradient shapes match the parameter shapes:

W1.grad: [2, 16]
b1.grad: [16]
W2.grad: [16, 2]
b2.grad: [2]

Autograd has already done the chain rule through:

cross entropy
    ↓
output linear layer
    ↓
ReLU
    ↓
first linear layer

We wrote the forward computation.

PyTorch derived the backward computation.

10. Update the parameters manually

Normally an optimizer does this.

For stochastic gradient descent:

parameter = parameter - learning_rate × gradient

In code:

learning_rate = 0.1

with torch.no_grad():
    W1 -= learning_rate * W1.grad
    b1 -= learning_rate * b1.grad
    W2 -= learning_rate * W2.grad
    b2 -= learning_rate * b2.grad

Why torch.no_grad()?

Because parameter updates are not part of the model we want to differentiate through.

If gradient tracking remained enabled, PyTorch would start recording the update itself as part of a new graph.

11. Zero the gradients

Gradients accumulate in PyTorch.

So after updating:

W1.grad.zero_()
b1.grad.zero_()
W2.grad.zero_()
b2.grad.zero_()

If you forget this, the next backward pass adds new gradients to the existing ones.

A helper keeps things cleaner:

parameters = [W1, b1, W2, b2]


def zero_grad(parameters):
    for p in parameters:
        if p.grad is not None:
            p.grad.zero_()

Then:

zero_grad(parameters)

12. Write the complete full-batch training loop

We now know every moving part.

import torch


torch.manual_seed(42)

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

indices = torch.randperm(n)
train_size = int(n * 0.8)

X_train = X[indices[:train_size]]
y_train = y[indices[:train_size]]
X_val = X[indices[train_size:]]
y_val = y[indices[train_size:]]

# parameters
W1 = (torch.randn(2, 16) * 0.1).requires_grad_()
b1 = torch.zeros(16, requires_grad=True)
W2 = (torch.randn(16, 2) * 0.1).requires_grad_()
b2 = torch.zeros(2, requires_grad=True)

parameters = [W1, b1, W2, b2]


def relu(x):
    return torch.clamp_min(x, 0.0)


def forward(x):
    hidden = relu(x @ W1 + b1)
    return hidden @ 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):
    predictions = logits.argmax(dim=1)
    return (predictions == targets).float().mean()


def zero_grad():
    for p in parameters:
        if p.grad is not None:
            p.grad.zero_()


learning_rate = 0.1
steps = 1000

# optional: track loss for a plot
loss_history = []

for step in range(steps):
    logits = forward(X_train)
    loss = cross_entropy(logits, y_train)

    zero_grad()
    loss.backward()

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

    loss_history.append(loss.item())

    if step % 100 == 0:
        with torch.no_grad():
            train_logits = forward(X_train)
            val_logits = forward(X_val)

            train_acc = accuracy(train_logits, y_train)
            val_acc = accuracy(val_logits, y_val)

        print(
            f"step={step:04d} "
            f"loss={loss.item():.4f} "
            f"train_acc={train_acc.item():.3f} "
            f"val_acc={val_acc.item():.3f}"
        )

You should see the loss decrease and classification accuracy climb substantially above chance.

The exact numbers vary with initialization and data.

The important thing is that this model learns a nonlinear decision boundary using nothing more than tensor operations and autograd.

Let’s see the loss curve:

import matplotlib.pyplot as plt

plt.plot(loss_history)
plt.xlabel('Step')
plt.ylabel('Loss')
plt.title('Manual training loss (full batch)')
plt.grid(True)
plt.show()

A steady decline means our hand‑built training loop works.

The training loop, stripped to its skeleton

Once the code is working, reduce it mentally to this:

for step in range(steps):
    logits = forward(X)
    loss = loss_fn(logits, y)

    zero_grad()
    loss.backward()

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

That is the core of training.

Compare it with idiomatic PyTorch later:

for X, y in loader:
    logits = model(X)
    loss = loss_fn(logits, y)

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

The abstractions remove bookkeeping.

They do not replace the mechanism.

13. Add mini-batches manually

Real training rarely performs every update using the entire dataset.

Let’s write a tiny batch iterator.

def batches(X, y, batch_size, shuffle=True):
    n = X.shape[0]

    if shuffle:
        indices = torch.randperm(n)
    else:
        indices = torch.arange(n)

    for start in range(0, n, batch_size):
        batch_idx = indices[start:start + batch_size]
        yield X[batch_idx], y[batch_idx]

Use it:

for xb, yb in batches(X_train, y_train, batch_size=64):
    print(xb.shape, yb.shape)
    break

Output:

torch.Size([64, 2]) torch.Size([64])

Now our training loop becomes:

learning_rate = 0.05
batch_size = 64
epochs = 50

for epoch in range(epochs):
    for xb, yb in batches(X_train, y_train, batch_size=batch_size):
        logits = forward(xb)
        loss = cross_entropy(logits, yb)

        zero_grad()
        loss.backward()

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

    if epoch % 5 == 0:
        with torch.no_grad():
            val_logits = forward(X_val)
            val_loss = cross_entropy(val_logits, y_val)
            val_acc = accuracy(val_logits, y_val)

        print(
            f"epoch={epoch:03d} "
            f"val_loss={val_loss.item():.4f} "
            f"val_acc={val_acc.item():.3f}"
        )

Congratulations.

You have now implemented the conceptual job of a DataLoader as well.

Not all of its functionality, obviously.

But the important training abstraction is visible:

    flowchart LR
    A[Shuffle indices] --> B[Slice batches]
    B --> C[Forward pass]
    C --> D[Compute loss]
    D --> E[Backward pass]
    E --> F[Update parameters]
    F --> B

14. Why initialization matters

Try this deliberately bad initialization:

W1 = torch.zeros(2, 16, requires_grad=True)
b1 = torch.zeros(16, requires_grad=True)
W2 = torch.zeros(16, 2, requires_grad=True)
b2 = torch.zeros(2, requires_grad=True)

Train the model.

It behaves very differently.

Why?

If neurons in a layer begin identically, they can receive identical gradients and remain symmetric.

Random initialization breaks that symmetry.

Now try ridiculously large initialization:

W1 = (torch.randn(2, 16) * 100).requires_grad_()

Inspect:

hidden_pre = X_train @ W1 + b1
hidden = relu(hidden_pre)

print(hidden_pre.abs().mean())
print(hidden.abs().mean())

Initialization influences signal scale, activation behavior and gradient scale.

Later, PyTorch layers will initialize their own weights according to defined schemes.

For now the important point is simple:

The initial values of the parameters are part of the training system.

15. Debug a model that is not learning

When the loss does not move, stop staring at the loss.

Inspect the system.

Are outputs changing?

with torch.no_grad():
    logits = forward(X_train[:8])
    print(logits)

Are gradients present?

loss = cross_entropy(forward(X_train), y_train)
zero_grad()
loss.backward()

for name, p in zip(["W1", "b1", "W2", "b2"], parameters):
    print(
        name,
        "grad_none=", p.grad is None,
        "grad_norm=", None if p.grad is None else p.grad.norm().item(),
    )

Are gradients finite?

for name, p in zip(["W1", "b1", "W2", "b2"], parameters):
    if p.grad is not None:
        print(name, torch.isfinite(p.grad).all().item())

Are parameters actually changing?

before = [p.detach().clone() for p in parameters]

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

for name, old, new in zip(
    ["W1", "b1", "W2", "b2"],
    before,
    parameters,
):
    delta = (new - old).abs().max().item()
    print(name, "max_update=", delta)

If gradients exist but parameters do not change, your update path is broken.

If parameters change but loss does not, investigate the data, loss, architecture and optimization scale.

If gradients are None, investigate the graph.

This is why building the machinery once by hand is useful.

You now know where to look.

16. Add a reusable parameter report

Programmers should instrument training code.

def parameter_report(named_parameters):
    rows = []

    for name, p in named_parameters:
        grad = p.grad

        rows.append({
            "name": name,
            "shape": tuple(p.shape),
            "param_norm": p.detach().norm().item(),
            "grad_none": grad is None,
            "grad_norm": None if grad is None else grad.detach().norm().item(),
            "grad_finite": None if grad is None else torch.isfinite(grad).all().item(),
        })

    return rows

Use it:

named_parameters = [
    ("W1", W1),
    ("b1", b1),
    ("W2", W2),
    ("b2", b2),
]

for row in parameter_report(named_parameters):
    print(row)

This becomes much more useful when your model has dozens of layers.

17. Device support without nn.Module

Because our model is just tensors, moving it to a device means moving every tensor ourselves.

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)

X_train = X_train.to(device)
y_train = y_train.to(device)
X_val = X_val.to(device)
y_val = y_val.to(device)

W1 = W1.to(device).detach().requires_grad_()
b1 = b1.to(device).detach().requires_grad_()
W2 = W2.to(device).detach().requires_grad_()
b2 = b2.to(device).detach().requires_grad_()

parameters = [W1, b1, W2, b2]

This is already annoying.

And we have only four parameters.

Imagine doing it for 400 tensors.

This is one of the problems nn.Module solves.

18. Saving without nn.Module

Again, because our model is just tensors, we need to define the state ourselves.

state = {
    "W1": W1.detach().cpu(),
    "b1": b1.detach().cpu(),
    "W2": W2.detach().cpu(),
    "b2": b2.detach().cpu(),
}

torch.save(state, "xor_manual.pt")

Load it:

state = torch.load("xor_manual.pt", map_location="cpu")

W1 = state["W1"].requires_grad_()
b1 = state["b1"].requires_grad_()
W2 = state["W2"].requires_grad_()
b2 = state["b2"].requires_grad_()

Again: possible, straightforward, repetitive.

Another reason nn.Module exists.

19. Count parameters manually

How many trainable numbers does our network contain?

for name, p in named_parameters:
    print(name, p.numel())

print("total:", sum(p.numel() for _, p in named_parameters))

The arithmetic:

W1: 2 × 16  = 32
b1: 16      = 16
W2: 16 × 2  = 32
b2: 2       = 2
-----------------
total          82

Our neural network has 82 trainable parameters.

A modern language model has the same conceptual object at vastly different scale:

collections of parameter tensors

20. Now rewrite it using nn.Module

Only now do we introduce the standard abstraction.

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


class XORNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(2, 16)
        self.fc2 = nn.Linear(16, 2)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return self.fc2(x)

Create it:

model = XORNet()
print(model)

Output looks roughly like:

XORNet(
  (fc1): Linear(in_features=2, out_features=16, bias=True)
  (fc2): Linear(in_features=16, out_features=2, bias=True)
)

Compare the implementations.

Manual:

hidden = relu(x @ W1 + b1)
logits = hidden @ W2 + b2

Module:

x = F.relu(self.fc1(x))
return self.fc2(x)

nn.Linear is not mysterious anymore.

Conceptually it owns a weight tensor and optional bias tensor and applies an affine transformation.

21. Inspect what nn.Module registered

This is one of the major benefits.

for name, p in model.named_parameters():
    print(name, tuple(p.shape), p.requires_grad)

You should see something like:

fc1.weight (16, 2) True
fc1.bias   (16,)   True
fc2.weight (2, 16) True
fc2.bias   (2,)    True

Notice something important.

Our manual matrix multiplication used:

W1: [2, 16]

PyTorch’s nn.Linear(2, 16) stores its weight as:

[16, 2]

The layer handles the transpose semantics internally.

This is exactly why reading shapes matters when you move between raw matrix code and layer APIs.

22. Verify an nn.Linear operation manually

Let’s prove it.

layer = nn.Linear(2, 16)
input_batch = torch.randn(8, 2)

layer_output = layer(input_batch)

manual_output = input_batch @ layer.weight.T + layer.bias

print(torch.allclose(layer_output, manual_output))

Output:

True

There is your nn.Linear abstraction.

At its core:

x @ weight.T + bias

23. Now use an optimizer

Our manual update was:

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

PyTorch packages optimization rules in torch.optim.

model = XORNet()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

Now:

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

replaces our bookkeeping.

The optimizer is not doing the backward pass.

loss.backward() computes gradients.

optimizer.step() reads those gradients and updates the parameters according to the optimizer’s rule.

That distinction matters when debugging.

24. Full idiomatic PyTorch version

Here is the same model using standard PyTorch abstractions.

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


torch.manual_seed(42)

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

indices = torch.randperm(n)
train_size = int(n * 0.8)

X_train = X[indices[:train_size]]
y_train = y[indices[:train_size]]
X_val = X[indices[train_size:]]
y_val = y[indices[train_size:]]


class XORNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(2, 16)
        self.fc2 = nn.Linear(16, 2)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return self.fc2(x)


model = XORNet()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

batch_size = 64
epochs = 50


def batches(X, y, batch_size):
    indices = torch.randperm(X.shape[0])

    for start in range(0, X.shape[0], batch_size):
        idx = indices[start:start + batch_size]
        yield X[idx], y[idx]


for epoch in range(epochs):
    model.train()

    for xb, yb in batches(X_train, y_train, batch_size):
        logits = model(xb)
        loss = F.cross_entropy(logits, yb)

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

    if epoch % 5 == 0:
        model.eval()

        with torch.no_grad():
            val_logits = model(X_val)
            val_loss = F.cross_entropy(val_logits, y_val)
            val_acc = (val_logits.argmax(dim=1) == y_val).float().mean()

        print(
            f"epoch={epoch:03d} "
            f"val_loss={val_loss.item():.4f} "
            f"val_acc={val_acc.item():.3f}"
        )

This is shorter.

More importantly, it composes.

The module owns its parameters.

The optimizer discovers them through:

model.parameters()

Device movement can be:

model.to(device)

Saving can use:

model.state_dict()

Nested modules register recursively.

This is what the abstraction buys us.

25. Manual vs nn.Module: what disappeared?

Here is the useful comparison.

Manual model state

parameters = [W1, b1, W2, b2]

nn.Module

model.parameters()

Manual forward

hidden = relu(x @ W1 + b1)
logits = hidden @ W2 + b2

nn.Module

hidden = F.relu(self.fc1(x))
logits = self.fc2(hidden)

Manual update

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

Optimizer

optimizer.step()

Manual gradient clearing

for p in parameters:
    p.grad.zero_()

Optimizer

optimizer.zero_grad()

Manual device movement

W1 = W1.to(device)
b1 = b1.to(device)
W2 = W2.to(device)
b2 = b2.to(device)

nn.Module

model.to(device)

The abstraction is mostly parameter registration and lifecycle management around tensor operations.

That is extremely valuable.

But it is no longer magic.

26. A subtle bug: replacing a parameter accidentally

When working manually, this mistake is easy:

W1 = W1 - learning_rate * W1.grad

Now you have assigned a new tensor to W1.

Depending on how you manage your parameter list, references elsewhere may still point at the old tensor.

Compare:

parameters = [W1]

old_id = id(W1)
W1 = W1 - 0.1 * W1.grad
new_id = id(W1)

print(old_id == new_id)
print(parameters[0] is W1)

This kind of bookkeeping problem is another reason optimizers update registered parameters in a controlled way.

When updating manually, prefer an in-place update inside torch.no_grad():

with torch.no_grad():
    W1 -= learning_rate * W1.grad

27. Another subtle bug: softmax before cross entropy

A common beginner implementation does this:

probs = torch.softmax(logits, dim=1)
loss = F.cross_entropy(probs, targets)

Do not do this.

F.cross_entropy expects logits and performs the relevant log-softmax / negative-log-likelihood computation internally.

Use:

loss = F.cross_entropy(logits, targets)

Apply softmax when you actually need probabilities for interpretation:

with torch.no_grad():
    probabilities = torch.softmax(logits, dim=1)

This distinction shows up constantly in real classifier code.

28. Check the model with explicit shape assertions

Treat shapes like runtime contracts.

def forward_checked(x):
    assert x.ndim == 2
    assert x.shape[1] == 2

    z1 = x @ W1 + b1
    assert z1.shape == (x.shape[0], 16)

    h = relu(z1)
    assert h.shape == (x.shape[0], 16)

    logits = h @ W2 + b2
    assert logits.shape == (x.shape[0], 2)

    return logits

During development, these assertions can turn a mysterious error 40 operations downstream into an immediate failure at the violated assumption.

29. Inspect activations while training

A model can have finite gradients and still behave badly.

Inspect intermediate activations.

def forward_debug(x):
    z1 = x @ W1 + b1
    h = relu(z1)
    logits = h @ W2 + b2

    print("input")
    print("  shape:", tuple(x.shape))
    print("  mean: ", x.mean().item())
    print("  std:  ", x.std().item())

    print("z1")
    print("  mean: ", z1.mean().item())
    print("  std:  ", z1.std().item())

    print("hidden")
    print("  mean: ", h.mean().item())
    print("  std:  ", h.std().item())
    print("  zeros:", (h == 0).float().mean().item())

    print("logits")
    print("  mean: ", logits.mean().item())
    print("  std:  ", logits.std().item())

    return logits

This style of instrumentation becomes useful when debugging dead activations, exploding values, normalization problems and unstable deep networks.

30. What nn.Module actually gives you

At this point it should be clear why PyTorch has nn.Module.

It gives structure to collections of tensors that belong to a model.

In particular, modules make it possible to:

model.parameters()
model.named_parameters()
model.state_dict()
model.to(device)
model.train()
model.eval()

Submodules are registered recursively, so a large model can be composed from smaller modules while still behaving as one object.

The underlying forward computation remains tensor operations.

A useful mental model

Think of a PyTorch model as three layers of machinery.

    flowchart TB
    subgraph Layer1["Layer 1: Tensor Math"]
        direction LR
        T1["x @ weight.T + bias"] --> T2["relu()"]
        T2 --> T3["softmax / log_softmax"]
    end
    subgraph Layer2["Layer 2: Autograd"]
        A1["Records operations"] --> A2["Computes gradients"]
        A2 --> A3["Populates .grad"]
    end
    subgraph Layer3["Layer 3: Abstractions"]
        M1["nn.Module / nn.Parameter"] --> M2["Optimizers"]
        M2 --> M3["DataLoader / state_dict"]
    end
    Layer1 --> Layer2 --> Layer3

We built Layers 1 and 2 ourselves.

Now Layer 3 has a reason to exist.

Challenge 1: add another hidden layer

Change:

2 → 16 → 2

to:

2 → 32 → 16 → 2

Create:

W1
b1
W2
b2
W3
b3

Write the entire forward pass manually.

Before running it, write down the shape of every intermediate tensor.

Challenge 2: implement sigmoid

Write:

def sigmoid(x):
    ...

using basic tensor operations.

Compare it against:

torch.sigmoid(x)

with:

torch.allclose(...)

Then replace ReLU in the network with sigmoid and compare training behavior.

Challenge 3: implement mean squared error

For a regression problem:

def mse(prediction, target):
    ...

Then verify against:

F.mse_loss(prediction, target)

Challenge 4: break the optimizer path

In the nn.Module version, deliberately construct an optimizer that only contains one layer:

optimizer = torch.optim.SGD(
    model.fc2.parameters(),
    lr=0.1,
)

Train it.

Then inspect:

for name, p in model.named_parameters():
    print(name, p.grad is None)

You may discover something important:

A parameter can receive a gradient and still not be updated if it is not registered with the optimizer.

This connects directly back to the debugging techniques in Step 02.

Challenge 5: verify nn.Linear numerically

For every linear layer in the model, prove that:

layer(x)

matches:

x @ layer.weight.T + layer.bias

Do not assume it.

Test it.

The full series so far

Step 00 — What Are We Actually Doing?
Step 01 — PyTorch Tensor Shapes and Broadcasting Errors
Step 02 — PyTorch Autograd Debugging
Step 03 — Build a Neural Network From Scratch Without nn.Module

We now understand:

tensors
    ↓
shapes
    ↓
operations
    ↓
autograd
    ↓
parameters
    ↓
forward pass
    ↓
loss
    ↓
backpropagation
    ↓
parameter updates
    ↓
mini-batch training

The next step is where we deliberately stop doing all the plumbing ourselves.

Next: PyTorch nn.Module Explained Through Real Code

In Step 04 we will focus on the abstractions programmers actually use every day:

nn.Module
nn.Parameter
state_dict
train() / eval()
hooks
module trees
parameter registration
buffers
optimizers

And because we are keeping the search-first approach, the next article will also tackle one of the most common PyTorch failure classes:

Why is my parameter missing from model.parameters() or state_dict()?

That problem makes much more sense now that we have built the underlying system ourselves.