Put it together: train and validate each epoch, track the loss curves, keep the best checkpoint, and save the model for inference. This is a real training script.
Why: training loss alone can fall while the model overfits — validating every epoch on held-out data is the only way to see real progress and catch overfitting. When: run a train pass then an eval pass each epoch, logging both losses. Where: switch train()/eval() modes around each phase, as covered earlier.
def run_epoch(model, loader, criterion, optimizer=None):
train = optimizer is not None
model.train() if train else model.eval()
total = 0.0
with torch.set_grad_enabled(train):
for x, y in loader:
pred = model(x)
loss = criterion(pred, y)
if train:
optimizer.zero_grad(); loss.backward(); optimizer.step()
total += loss.item() * len(x)
return total / len(loader.dataset)Why: the last epoch is rarely the best — you save the model whenever validation improves, so you keep the peak instead of an overfit final state. When: compare validation loss each epoch and checkpoint on improvement. Where: save state_dict (the weights), not the whole model object, for portable reloading.
best_val = float("inf")
for epoch in range(50):
train_loss = run_epoch(model, train_loader, criterion, optimizer)
val_loss = run_epoch(model, val_loader, criterion) # no optimizer
print(f"epoch {epoch} train {train_loss:.4f} val {val_loss:.4f}")
if val_loss < best_val: # improved -> checkpoint
best_val = val_loss
torch.save(model.state_dict(), "best.pt")Why: a trained model is only useful if you can reload it later for predictions — you save the weights and, to use them, rebuild the architecture and load them in. When: load the best checkpoint, switch to eval mode, and predict with no_grad. Where: you must recreate the same model class before loading its state_dict.
# Later, in a serving script:
model = MLP(in_features=10, hidden=32, out=3) # same architecture
model.load_state_dict(torch.load("best.pt"))
model.eval()
with torch.no_grad():
logits = model(new_x)
preds = logits.argmax(dim=1) # predicted class per sample
print(preds)