Training Loop

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = MyNet().to(device)
loss_fn   = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-2)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
scaler    = torch.amp.GradScaler('cuda')        # mixed precision
writer    = SummaryWriter('runs/exp1')

best_acc = 0.0
for epoch in range(1, epochs + 1):
    # ---- TRAIN ----
    model.train()
    running_loss = 0.0
    for step, (x, y) in enumerate(train_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_()

        with torch.amp.autocast('cuda', dtype=torch.float16):
            logits = model(x)
            loss   = loss_fn(logits, y)

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

        running_loss += loss.item() * x.size(0)
        if step % 50 == 0:
            writer.add_scalar('train/loss', loss.item(), epoch * len(train_loader) + step)

    scheduler.step()
    train_loss = running_loss / len(train_loader.dataset)

    # ---- VALIDATE ----
    model.eval()
    correct, total = 0, 0
    with torch.inference_mode():
        for x, y in val_loader:
            x, y = x.to(device), y.to(device)
            logits = model(x)
            preds  = logits.argmax(dim=1)
            correct += (preds == y).sum().item()
            total   += y.size(0)
    val_acc = correct / total

    writer.add_scalar('train/epoch_loss', train_loss, epoch)
    writer.add_scalar('val/acc', val_acc, epoch)
    print(f'ep {epoch} train_loss={train_loss:.4f} val_acc={val_acc:.4f}')

    # ---- CHECKPOINT THE BEST ----
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save({
            'model':     model.state_dict(),
            'optim':     optimizer.state_dict(),
            'scaler':    scaler.state_dict(),
            'epoch':     epoch,
            'val_acc':   val_acc,
        }, 'best.pt')

# Load later
ckpt = torch.load('best.pt', weights_only=True)
model.load_state_dict(ckpt['model'])

model.train()
for xb, yb in train_loader:
    xb, yb = xb.to(dev), yb.to(dev)
    pred = model(xb)
    loss = loss_fn(pred, yb)
    opt.zero_grad()
    loss.backward()
    opt.step()