for epoch in ... that actually trains a model. You'll meet loss functions, optimizers (Stochastic Gradient Descent, or SGD, and Adam), the learning rate, and how to catch overfitting before it ruins you.Whether it's a tiny classifier or GPT-4, training is this loop, repeated. Five steps, over and over, each pass making the model slightly less wrong:
for epoch in range(num_epochs):
pred = model(x) # 1. FORWARD: make a prediction
loss = loss_fn(pred, y) # 2. LOSS: how wrong are we?
optimizer.zero_grad() # 3. reset old gradients
loss.backward() # 4. BACKWARD: compute gradients (autograd!)
optimizer.step() # 5. UPDATE: nudge weights downhillStep 4 (backward) is autograd from Module 2. Step 5 (step) is x = x − lr×grad, now handled by an optimizer object so you don't write it by hand. Steps 1–3 are new but simple: predict, measure error, clear the old slopes. Master these 5 lines and you can train anything.
Click each step to see what it does. This is the exact order every single training run follows.
This trains a tiny model to fit a line. Press "Train" and watch the loss curve fall as the model learns — the payoff of the whole loop. (Simulated to show the shape; the real code is below.)
import torch, torch.nn as nn
torch.manual_seed(1) # same starting weight each run
x = torch.tensor([[1.0], [2.0], [3.0], [4.0]]) # inputs: shape (4, 1)
y = torch.tensor([[3.0], [5.0], [7.0], [9.0]]) # targets: y = 2x + 1
model = nn.Linear(1, 1) # y = w*x + b
loss_fn = nn.MSELoss() # mean squared error
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for epoch in range(10):
pred = model(x)
loss = loss_fn(pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(epoch, loss.item())A loss function measures "how wrong." Different jobs need different rulers. The two you'll use 90% of the time:
The optimizer decides how to use the gradient to update weights. Stochastic Gradient Descent (SGD) takes plain steps. Adam is smarter — it adapts the step size per-weight and adds "momentum," so it usually converges faster. Watch them race down the same loss.
# swapping optimizers is one line:
optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # Stochastic Gradient Descent
optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # adaptive, popular defaultThe learning rate is the single most important setting. Slide it and run — see the loss curve go smooth, sluggish, or explode. This is the intuition you'll rely on constantly.
Here's the trap: a model can get a perfect score on its training data by memorizing it — then fail completely on anything new. That's overfitting. To catch it, you split your data:
Train set (~80%) — the model learns from this. Validation set (~10%) — checked during training to see if it generalizes (and tune settings). Test set (~10%) — touched once at the very end, for an honest final grade. The rule: never let the model train on your test set.
Train = the practice problems you study. Validation = a mock exam to check you actually understand (not just memorized answers). Test = the real exam, seen once. A student who memorizes the practice answers aces practice but bombs the real test — that's overfitting.
Train loss keeps dropping — but watch the validation loss. When val starts rising while train keeps falling, the model has stopped learning and started memorizing. Run it and find that moment.
When validation loss stops improving (or rises), stop training — that's called early stopping. It's the simplest overfitting cure. Two others come up constantly: dropout and train/eval mode — let's cover them now.
train() / eval() switchDropout randomly "switches off" a fraction of neurons on each training step. It sounds destructive, but it forces the network to not over-rely on any single neuron — so it generalizes better instead of memorizing. It's one of the most common regularization tools.
If the same one person always answers, the group learns to lean on them. Randomly benching members each session forces everyone to actually learn the material — so the group is robust even when someone's missing. Dropout does that to neurons.
model.train() vs model.eval()Here's a bug that bites almost every beginner. Dropout should be ON while training (to regularize) but OFF when evaluating or deploying (you want the full, stable network — no random neurons missing). PyTorch controls this with a mode switch you must set yourself:
model = nn.Sequential(
nn.Linear(10, 64), nn.ReLU(),
nn.Dropout(0.2), # randomly drop 20% of neurons
nn.Linear(64, 1),
)
# --- during TRAINING ---
model.train() # dropout ON (regularizes)
for batch in train_loader:
... # the 5-step loop
# --- during VALIDATION / TEST / INFERENCE ---
model.eval() # dropout OFF (full, deterministic network)
with torch.no_grad(): # also skip gradient tracking → faster, less memory
preds = model(val_x)If you evaluate while still in train() mode, dropout stays on and your validation scores come out randomly worse (and different every run). If you train while in eval() mode, you lose your regularization. The habit: model.train() before the training loop, model.eval() before you measure or predict. Pair eval() with torch.no_grad() so you're not wasting memory tracking gradients you won't use.
Some layers behave differently in training vs. inference — Dropout (random off vs. all on) and BatchNorm (uses batch stats vs. running averages) are the big two. Rather than configure each layer, you flip the whole model's mode with one call, and every layer adjusts itself. That's what train()/eval() do.
forward → loss → zero_grad → backward → step. A loss function measures wrongness (MSE for numbers, Cross-Entropy for categories). An optimizer (Stochastic Gradient Descent, or SGD, or the popular Adam) turns gradients into weight updates. The learning rate is the most important knob — too small crawls, too big explodes. Split data into train / validation / test; when validation loss rises while train falls, you're overfitting — stop early.Next up → Module 4: Datasets & Tokenizers — where the x and y in this loop actually come from. Loading real data with Hugging Face, batching, and turning text into tokens the model can eat.