PyTorch Training Loop

The PyTorch training loop is code you write yourself, which is one of the things people love about PyTorch: nothing is hidden. Every iteration follows the same five steps, and once you know them you can train almost any model.

The Five Steps

• 1. Forward pass: pred = model(x)
• 2. Compute the loss: loss = loss_fn(pred, y)
• 3. Clear old gradients: optimizer.zero_grad()
• 4. Backpropagate: loss.backward()
• 5. Update weights: optimizer.step()

A Complete Example: Learning a Line

We generate data from y = 3x + 2 with some noise and let a single linear layer discover the slope and intercept.

import torch
from torch import nn

torch.manual_seed(42)
X = torch.rand(200, 1) * 10
y = 3 * X + 2 + torch.randn(200, 1) * 0.5

model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
opt = torch.optim.SGD(model.parameters(), lr=0.01)

for epoch in range(1000):
    pred = model(X)
    loss = loss_fn(pred, y)
    opt.zero_grad()
    loss.backward()
    opt.step()
    if epoch % 250 == 0:
        print(f"epoch {epoch}: loss {loss.item():.4f}")

print(model.weight.item(), model.bias.item())

epoch 0: loss 160.3521
epoch 250: loss 0.7012
epoch 500: loss 0.3874
epoch 750: loss 0.2789
2.97... 2.18...
(sample output — values vary slightly)

The learned weight is close to 3 and the bias close to 2. That is the whole idea of training: adjust parameters to reduce the loss.

Adding a Validation Loop

Always measure performance on data the model did not train on. Evaluation uses model.eval() and torch.no_grad():

def evaluate(model, X_val, y_val):
    model.eval()
    with torch.no_grad():
        loss = loss_fn(model(X_val), y_val).item()
    model.train()
    return loss

X_val = torch.rand(50, 1) * 10
y_val = 3 * X_val + 2
print(round(evaluate(model, X_val, y_val), 4))

0.0412   (sample output)

Tracking Accuracy for Classification

logits = torch.tensor([[2.0, 0.1], [0.3, 1.2], [1.0, 0.9]])
labels = torch.tensor([0, 1, 1])
preds = logits.argmax(dim=1)
acc = (preds == labels).float().mean().item()
print(preds, acc)

tensor([0, 1, 0]) 0.6666666865348877

Reusable Train Function

Real projects loop over mini-batches from a DataLoader (next lesson). A reusable function keeps the loop tidy:

def train_one_epoch(model, loader, loss_fn, opt, device):
    model.train()
    total = 0.0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        loss = loss_fn(model(xb), yb)
        opt.zero_grad()
        loss.backward()
        opt.step()
        total += loss.item() * xb.size(0)
    return total / len(loader.dataset)

(returns the average loss for the epoch)

Frequently asked questions

What is the correct order of zero_grad, backward and step?

Call optimizer.zero_grad() before loss.backward(), then optimizer.step(). zero_grad can also go at the start of the iteration; it just must happen before backward.

What is an epoch in PyTorch?

One complete pass over the entire training dataset. Training usually runs for several epochs.

Why does my loss become NaN?

Common causes are a learning rate that is too high, unnormalized inputs, taking log of zero, or bad data. Lower the learning rate and check your inputs first.

Does PyTorch have a fit() function like Keras?

Not in core PyTorch; you write the loop. Libraries such as PyTorch Lightning or Hugging Face Trainer provide fit-style APIs on top.