Level 1 · FoundationsStage 03

Lesson 11 of 38 · Written learning

Training a small model with PyTorch

Connect tensors, automatic differentiation, batches, and optimizers in a training loop.

6 min read Free course lessonWeeks 6–8

PyTorch tensors hold numerical data and can track operations for automatic differentiation. A batch is a small group of examples processed together. Batching uses hardware efficiently and gives the optimizer an estimate of the direction that reduces loss across the batch.

The training loop in plain language

First, the model produces logits. A loss compares those logits with target values. Backpropagation computes gradients, and an optimizer updates parameters. Before the next batch, clear old gradients so they are not accidentally accumulated. Evaluate on validation data without updating parameters.

Illustrative PyTorch pattern
optimizer.zero_grad()
logits = model(inputs)
loss = loss_fn(logits, targets)
loss.backward()
optimizer.step()
  • Track training and validation loss separately.
  • Save checkpoints when an experiment is reproducible and useful to resume.
  • A GPU can accelerate suitable operations, but small examples can be learned on a CPU.
  • If training loss falls while validation loss rises, investigate overfitting, data splits, and model capacity.

Return to the Machine learning and deep learning stage · Suggested timing is flexible; move at your own pace.