PyTorch Training Loop Explainer
Interactive visual guide to tensors, gradients, and how neural networks learn
Animated Computation Graph: loss.backward()
Watch how gradients flow backward through the computation graph when loss.backward() is called.
Click "Run loss.backward()" to see gradient flow animation.
Tensor & Matmul Shape Diagrams
Visualizing matrix multiplication shapes in a neural network layer.
Input X
[1, 4]
batch x features
Weight W
[4, 3]
features x output
Output = X @ W
[1, 3]
batch x output
# Matrix multiplication: output = input @ weight
# Shapes: [batch, features] @ [features, output] -> [batch, output]
# Inner dimensions must match (features = features)
Step-Through Mini Training Loop
Watch weights update live as we train a simple linear model.
predictions = model(inputs)
loss = loss_fn(predictions, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
Weight: 0.500
Bias: 0.100
Gradient: 0.000
Quiz: Test Your Understanding
Answer questions about each line of the training loop.
1. What does loss.backward() compute?
The forward pass output
Gradients of loss w.r.t. all parameters
The learning rate schedule
The training accuracy
2. Why do we call optimizer.zero_grad() before loss.backward()?
To reset the model weights
To clear accumulated gradients from previous steps
To initialize the optimizer
To free GPU memory
3. What does optimizer.step() do?
Computes the loss
Updates parameters using gradients
Saves the model checkpoint
Evaluates on test data
4. In matrix multiplication [batch, features] @ [features, output], what must match?
batch and output
features and features (inner dimensions)
batch and features
output and batch