PyTorch: an honest training loop

The previous lesson traced one error backwards by hand. That is valuable because it tells you what a gradient means. But manually writing derivatives for every parameter stops being practical as soon as the model has more than a few calculations.

PyTorch is a Python library for numerical computation and deep learning. A tensorTensorA container for numbers arranged with one or more dimensions. A single number is a zero-dimensional tensor, a list is often one-dimensional, and a table is often two-dimensional. In PyTorch, tensors can hold data and, when needed, record operations for gradient calculation. stores its numbers; the library can record eligible calculations for automatic differentiation and provides model and optimiser building blocks. It does not choose a valid target, protect a data split, or decide that a model is safe to use.

Why this matters

A training framework makes the repeated arithmetic manageable. It also makes it easy to hide mistakes behind five short method calls. A readable training loop gives every call one job: make a prediction, measure the error, clear old gradients, calculate new gradients, then update parameters.

Plain-English mental model

Step 1 — Tensors are shaped containers for numbers

A tensor is a container whose shape says how numbers are arranged. A score for one example might have shape (1,). A batch of four examples, each with three input features, has shape (4, 3): four rows and three values per row.

PyTorch can run the same tensor code on a CPU or a supported accelerator. Start on the CPU, make shapes and types clear, and move only after measurement gives you a reason. Inputs for a neural network are usually floating-point tensors. Classification targets often have a different type, depending on the loss; always read the loss function’s contract rather than guessing.

import torch

features = torch.tensor(
    [[0.0], [1.0], [2.0]],
    dtype=torch.float32,
)
targets = torch.tensor([[1.0], [3.0], [5.0]], dtype=torch.float32)

print(features.shape)
print(features.dtype)
print(targets.shape)

The expected output is:

torch.Size([3, 1])
torch.float32
torch.Size([3, 1])

The outer dimension is the example or batch dimension. Keeping it—even for one example—helps a model accept one example and many examples consistently.

Step 2 — Keep the five training jobs separate

Here is the essential order. It is short enough to memorise, but each line has a distinct responsibility.

prediction = model(features)               # 1. forward pass
loss = loss_fn(prediction, targets)        # 2. numeric training signal

optimizer.zero_grad()                      # 3. remove old accumulated gradients
loss.backward()                            # 4. calculate fresh gradients
optimizer.step()                           # 5. update parameters

PyTorch adds gradients to a parameter’s .grad field by default. That is useful when you intentionally combine several smaller computations before an update. In the ordinary one-batch pattern, forgetting zero_grad() means a later update mixes in an earlier batch’s gradient. Clear before backward(), so the gradient you inspect belongs to the current batch.

Worked example — learn a line with an explicit model

This intentionally tiny problem has the rule target = 2 × input + 1. It is not a realistic product task. It lets us inspect the mechanics without pretending a short example proves generalisation.

nn.Linear(1, 1) is a model with one input, one output, a weight, and a bias. We set both parameters to zero so this first update is repeatable and understandable. The mean-squared-error loss averages the squared prediction errors.

import torch
from torch import nn

features = torch.tensor([[0.0], [1.0]])
targets = torch.tensor([[1.0], [3.0]])  # target = 2 × input + 1

model = nn.Linear(1, 1)
with torch.no_grad():
    model.weight.fill_(0.0)
    model.bias.fill_(0.0)

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

prediction = model(features)
loss = loss_fn(prediction, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()

print('loss before:', round(loss.item(), 3))
print('weight after:', round(model.weight.item(), 3))
print('bias after:', round(model.bias.item(), 3))

The result does not jump to weight 2 and bias 1 after one batch. One small update moves in a helpful direction. Training repeats on data that represents the intended task, while validation checks whether progress also appears on protected examples.

Step 3 — A DataLoader provides batches, not validity

A DataLoaderDataLoaderA PyTorch object that iterates over a dataset and yields batches. It can batch examples, choose an order, and optionally load data using worker processes. It does not make a split valid or a dataset representative; those are separate design decisions. turns a dataset into an iterator over batches. It can group examples, shuffle a training order, and use worker processes when reading data becomes a measured bottleneck. It does not create a train/validation/test split for you, prevent leakage, make labels correct, or make shuffling appropriate for time-ordered data.

from torch.utils.data import DataLoader, TensorDataset

dataset = TensorDataset(features, targets)
train_loader = DataLoader(dataset, batch_size=2, shuffle=True)

for batch_features, batch_targets in train_loader:
    print(batch_features.shape, batch_targets.shape)

For this two-example dataset, the output shapes are always torch.Size([2, 1]) and torch.Size([2, 1]); only the row order may differ because shuffle=True changes it. Set shuffle=True only for a training split when the examples may safely be reordered. Do not casually shuffle a time sequence, grouped records, or a validation set merely because the default training pattern did.

An epochEpochOne pass through the training examples as defined by a data loader. An epoch is an accounting unit, not proof that a model has learned enough or will work on new data. means one pass through the examples yielded by a loader. It is a convenient counter, not a target by itself. Record both the training loss and a protected validation measure at each epoch so you can see divergence, overfitting, or a plateau rather than relying on a final number alone.

Step 4 — Train and evaluate deliberately differently

The training step needs gradients. Evaluation does not. During evaluation, use model.eval() and torch.no_grad().

model.eval() switches certain layers—such as dropout and batch normalisation—to their evaluation behaviour. It does not disable gradients. torch.no_grad() stops PyTorch recording gradient-tracking operations for the enclosed block. It does not switch model behaviour. In real code, use both for ordinary validation or inference.

for epoch in range(20):
    model.train()
    for batch_features, batch_targets in train_loader:
        prediction = model(batch_features)
        loss = loss_fn(prediction, batch_targets)

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

    model.eval()
    with torch.no_grad():
        validation_prediction = model(validation_features)
        validation_loss = loss_fn(validation_prediction, validation_targets)

    print(epoch, validation_loss.item())

Keep validation examples separate before you create the loader. A clean-looking loop cannot repair a contaminated split. For a real project, make preprocessing an explicit, versioned part of the training and evaluation path; fit data-dependent transformations on training data only, then apply the frozen transformation elsewhere.

Try it — see why accumulated gradients matter

This browser exercise uses plain Python to make the accumulation rule visible. The arithmetic matches a simplified one-parameter update; PyTorch performs the same bookkeeping for tensors. Run it, then change clear_before_second_batch to False and compare the final weight.

With clear_before_second_batch = True, the output is:

after batch 1: 0.4
gradient used for batch 2: -2.0
final weight: 0.6000000000000001

If you set it to False, the second update uses -6.0 instead: the old -4.0 has been added to the new -2.0. Gradient accumulation can be intentional, but it must be an explicit design choice with the loss scaling checked carefully.

Step 5 — Save model state, context, and evidence

A model stateModel stateThe values a model needs to reproduce its learned behaviour, including its parameters and persistent buffers. A PyTorch module exposes this information through state_dict; a usable training checkpoint commonly also records optimiser state and run metadata. is the learned parameter and persistent-buffer values needed for a module to reproduce its behaviour. In PyTorch, model.state_dict() provides that mapping. If you may resume training, store optimizer.state_dict() too; optimisers such as Adam keep internal history that affects future updates.

checkpoint = {
    'model_state': model.state_dict(),
    'optimizer_state': optimizer.state_dict(),
    'epoch': epoch,
    'validation_loss': validation_loss.item(),
    'data_version': 'support-labels-2026-07-15',
}
torch.save(checkpoint, 'checkpoint.pt')

The filename and one loss value are not enough evidence. Record the code revision, package versions, model configuration, data and split version, preprocessing version, random seeds where applicable, metrics and error slices, and the decision that determined whether this run may proceed. Never place secrets or unnecessary personal data into a checkpoint or its logs.

Real-world AI engineering example — a classifier that can be investigated

Common mistakes

Knowledge check

Practice activity

Key takeaways

Flashcards

One-page sketchnote summary

predict → loss
clear → backward → step
eval: eval() + no_grad()

Resources