Learn / PyTorch
Advanced

Loss Functions & Training Loop

Measuring how wrong a model is, and the update loop that fixes it.

A training step is always the same four moves: zero old gradients, forward pass to get predictions and loss, backward to compute gradients, step the optimizer to update parameters using those gradients. Running that loop once over the whole training set is one .

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(epochs):
    optimizer.zero_grad()          # 1. zero
    preds = model(X_train)          # 2. forward
    loss = loss_fn(preds, y_train)
    loss.backward()                 # 3. backward
    optimizer.step()                # 4. step
📝 Note: nn.CrossEntropyLoss expects raw logits and integer class labels; it applies log-softmax internally. Don't apply softmax yourself first, or you'll double-apply it.

Now practice it

Exercises that use what this lesson just covered.

Building Models with nn.ModuleClassification Metrics & Putting It Together