08_training_loop

08 — A reusable training loop

The patterns in this file show up in nearly every PyTorch project worth its salt.

What's in here

• Device handling — model.to(device) and xb.to(device) keep CPU/GPU code uniform.
• run_epoch helper — one function for both train and eval, controlled by a train: bool flag and torch.set_grad_enabled.
• Metrics — track loss and accuracy per epoch on both splits.
• Early stopping — if val loss hasn't improved in patience epochs, stop. Saves time and prevents overfitting.
• Checkpointing — torch.save({"model": model.state_dict(), ...}) writes weights + metadata. load_state_dict restores them.
• TrainConfig dataclass — keep hyperparameters in one place; easy to log or sweep.

State dict, not the model

Save model.state_dict(), not model itself. The state dict is a plain {name: tensor} mapping — portable, version-agnostic, and doesn't depend on the exact class definition being importable. Pickled nn.Module objects break on refactors.

weights_only=True

torch.load(..., weights_only=True) is the safe loader. The default unpickler can execute arbitrary code from a malicious checkpoint. Use weights_only=True for anything you don't trust 100%.

Run

python 08_training_loop/main.py