PyTorch Training Loop: Forward, Loss, Backward, Optimizer and Validation

Loop Structure

The training loop is where PyTorch gives you full control. You decide how batches move to the device, how loss is computed, when gradients are cleared, whether gradients are clipped, how metrics are tracked, and when checkpoints are saved.

A clean loop separates training and validation. Training uses `model.train()` and gradients. Validation uses `model.eval()` and `torch.no_grad()`. Mixing those modes is a common source of unreliable metrics.

Each epoch processes every training batch, then evaluates on validation data. Metrics should be averaged by the number of examples, not by a naive number of batches when batch sizes vary.

Move batch tensors to device.
Forward pass through model.
Compute loss.
Zero gradients, backward pass, optional clipping, optimizer step.
Switch to eval mode for validation.

Reusable Train and Validate Functions

This structure is easy to test and extend with schedulers, mixed precision, early stopping, and logging.

Reusable Train and Validate Functions

import torch

def train_one_epoch(model, loader, loss_fn, optimizer, device):
    model.train()
    total_loss = 0.0
    total_correct = 0
    total_examples = 0

    for features, labels in loader:
        features = features.to(device)
        labels = labels.to(device)

        logits = model(features)
        loss = loss_fn(logits, labels)

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

        batch_size = labels.size(0)
        total_loss += loss.item() * batch_size
        total_correct += (logits.argmax(dim=1) == labels).sum().item()
        total_examples += batch_size

    return {
        "loss": total_loss / total_examples,
        "accuracy": total_correct / total_examples,
    }

@torch.no_grad()
def validate(model, loader, loss_fn, device):
    model.eval()
    total_loss = 0.0
    total_correct = 0
    total_examples = 0

    for features, labels in loader:
        features = features.to(device)
        labels = labels.to(device)
        logits = model(features)
        loss = loss_fn(logits, labels)

        batch_size = labels.size(0)
        total_loss += loss.item() * batch_size
        total_correct += (logits.argmax(dim=1) == labels).sum().item()
        total_examples += batch_size

    return {
        "loss": total_loss / total_examples,
        "accuracy": total_correct / total_examples,
    }

Gradient clipping can prevent unstable updates in some models.
The validation function is decorated with no_grad to save memory and computation.

Run One Complete Optimization Step

import torch

model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
x = torch.tensor([[1.0], [2.0]])
y = torch.tensor([[2.0], [4.0]])

optimizer.zero_grad()
loss = torch.nn.functional.mse_loss(model(x), y)
loss.backward()
optimizer.step()

print(loss.ndim, all(parameter.grad is not None for parameter in model.parameters()))

Output

0 True

Before you move on