Why do we need to call zero_grad in PyTorch?
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
Introduction
In PyTorch, gradients are not overwritten automatically on each backward pass. They accumulate in each parameter’s .grad buffer, which is why a normal training loop calls zero_grad() before computing the next batch.
Why Gradients Accumulate
When you call loss.backward(), PyTorch computes gradients and adds them to existing values in param.grad. This is useful because it lets you accumulate gradients across multiple backward passes on purpose. But it also means that if you forget to clear them, the next optimization step uses the sum of old and new gradients.
That behavior is different from what many beginners expect. They assume each backward pass starts from a clean slate. In PyTorch, it does not.
A Typical Training Loop
The standard pattern is:
- Clear previous gradients.
- Run the forward pass.
- Compute the loss.
- Run backpropagation.
- Update the parameters.
If you remove optimizer.zero_grad(), the next iteration will add new gradients on top of the previous ones.
What Goes Wrong If You Skip It
Suppose your first backward pass produces gradient g1 and the second produces g2. If you do not clear gradients between them, PyTorch stores g1 + g2 in the gradient buffer. The optimizer then updates parameters using the accumulated value.
That usually means the model trains incorrectly, because each batch is no longer contributing exactly once.
This accumulation is not a bug. It is the designed behavior.
When Accumulation Is Actually Useful
Sometimes you intentionally want to accumulate gradients. A common reason is memory pressure. Instead of using one very large batch, you can process several smaller batches, call backward() on each one, and update the optimizer only after a few steps.
The division by accumulation_steps keeps the effective gradient scale close to what you would get from one larger batch.
model.zero_grad() vs optimizer.zero_grad()
You will see both forms in real code. optimizer.zero_grad() is usually the clearer choice because it resets exactly the parameters managed by that optimizer. model.zero_grad() also works when all trainable parameters belong to the model and the optimizer covers them all.
Modern PyTorch also supports optimizer.zero_grad(set_to_none=True). Setting gradients to None instead of filling them with zeros can save memory and slightly improve performance in some workloads.
Common Pitfalls
- Forgetting to clear gradients between batches is the classic training-loop bug.
- Accumulating gradients intentionally without scaling the loss changes the effective learning rate.
- Calling
zero_grad()afterbackward()but before inspecting gradients can confuse debugging. - Mixing multiple optimizers requires care, because each optimizer manages its own parameter set.
- Assuming gradients reset automatically because other frameworks do it differently leads to subtle mistakes.
Summary
- PyTorch accumulates gradients in
.gradbuffers by default. - '
zero_grad()clears those buffers so each batch starts clean.' - Skipping it usually causes incorrect parameter updates.
- Gradient accumulation is a deliberate technique, not something you want by accident.
- Use
optimizer.zero_grad()as the standard reset point in a training loop.
Related reading
- Why do we use fully-connected layer at the end of CNN?
- Why do we use tf.name_scope
- Why does a binary Keras CNN always predict 1?
- Why does a binary Keras CNN always predict 1?
- Why doesn't my simple pytorch network work on GPU device?
- Why is PyTorch 2x slower than Keras for an identical model and hyperparameters?
- Why does embedding vector multiplied by a constant in Transformer model?
- Why does my keras LSTM model get stuck in an infinite loop?
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack what you have practised
A free account saves your progress, solutions and study plan across every problem on Codemia.
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.