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

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}')