What does a production-quality PyTorch training loop look like, incorporating all best practices?

A well-structured training loop separates concerns cleanly: data loading, forward pass, loss computation, backpropagation, gradient management, metric tracking, and model persistence. Each step has specific pitfalls that silently degrade results.

import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler
from torch.utils.data import DataLoader

def train_epoch(model, loader, optimizer, criterion, device, scaler):
    model.train()
    total_loss, n_correct, n_total = 0.0, 0, 0

    for X, y in loader:
        X, y = X.to(device, non_blocking=True), y.to(device, non_blocking=True)

        optimizer.zero_grad(set_to_none=True)  # faster than zero_grad()

        with autocast(device_type='cuda', dtype=torch.float16):
            logits = model(X)
            loss   = criterion(logits, y)

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()

        total_loss += loss.item() * X.size(0)
        n_correct  += (logits.argmax(1) == y).sum().item()
        n_total    += X.size(0)

    return total_loss / n_total, n_correct / n_total

@torch.no_grad()
def eval_epoch(model, loader, criterion, device):
    model.eval()
    total_loss, n_correct, n_total = 0.0, 0, 0
    for X, y in loader:
        X, y = X.to(device, non_blocking=True), y.to(device, non_blocking=True)
        logits = model(X)
        loss   = criterion(logits, y)
        total_loss += loss.item() * X.size(0)
        n_correct  += (logits.argmax(1) == y).sum().item()
        n_total    += X.size(0)
    return total_loss / n_total, n_correct / n_total

# Main training loop
best_val_acc = 0
for epoch in range(n_epochs):
    tr_loss, tr_acc = train_epoch(model, train_loader, optimizer,
                                   criterion, device, scaler)
    vl_loss, vl_acc = eval_epoch(model, val_loader, criterion, device)
    scheduler.step()
    if vl_acc > best_val_acc:
        best_val_acc = vl_acc
        torch.save(model.state_dict(), 'best.pt')
    print(f'Epoch {epoch:3d}: tr={tr_loss:.4f}/{tr_acc:.3f}  '
          f'val={vl_loss:.4f}/{vl_acc:.3f}')