Why do we need to call zero_grad in PyTorch?
Master System Design with Codemia
Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.
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.

