How do you compute and track evaluation metrics like accuracy during PyTorch training?

import torch
import torch.nn as nn

@torch.no_grad()   # disable gradient tracking for the entire evaluation function
def evaluate(model, loader, criterion, device):
    model.eval()                       # disable dropout, use BN running stats

    total_loss    = 0.0
    total_correct = 0
    total_samples = 0

    for X_batch, y_batch in loader:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)
        batch_size = X_batch.size(0)

        logits = model(X_batch)
        loss   = criterion(logits, y_batch)

        # Weight by batch_size — correct even if the last batch is smaller
        total_loss += loss.item() * batch_size

        preds = logits.argmax(dim=1)
        total_correct += (preds == y_batch).sum().item()
        total_samples += batch_size

    avg_loss = total_loss / total_samples
    accuracy = total_correct / total_samples
    return avg_loss, accuracy

# WRONG pattern — naively averaging per-batch averages
# is biased if batch sizes are unequal (e.g. last batch is smaller)
def evaluate_wrong(model, loader, criterion):
    losses = []
    for X_batch, y_batch in loader:
        loss = criterion(model(X_batch), y_batch)
        losses.append(loss.item())   # all batches weighted EQUALLY — wrong!
    return sum(losses) / len(losses)  # biased if last batch has fewer samples

# Using torchmetrics for more complex metrics (F1, precision, AUROC)
# pip install torchmetrics
from torchmetrics import Accuracy, F1Score

acc_metric = Accuracy(task="multiclass", num_classes=5).to(device)
f1_metric  = F1Score(task="multiclass", num_classes=5, average="macro").to(device)

for X_batch, y_batch in loader:
    preds = model(X_batch).argmax(dim=1)
    acc_metric.update(preds, y_batch)   # accumulates internally across batches
    f1_metric.update(preds, y_batch)

print(f"Accuracy: {acc_metric.compute():.4f}")  # final correct aggregate
print(f"F1: {f1_metric.compute():.4f}")