PyTorch
zero_grad
deep learning
neural networks
backpropagation

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:

  1. Clear previous gradients.
  2. Run the forward pass.
  3. Compute the loss.
  4. Run backpropagation.
  5. Update the parameters.
python
1import torch
2import torch.nn as nn
3import torch.optim as optim
4
5model = nn.Linear(4, 1)
6optimizer = optim.SGD(model.parameters(), lr=0.1)
7criterion = nn.MSELoss()
8
9x = torch.randn(8, 4)
10y = torch.randn(8, 1)
11
12optimizer.zero_grad()
13pred = model(x)
14loss = criterion(pred, y)
15loss.backward()
16optimizer.step()

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.

python
1optimizer.zero_grad()
2loss1.backward()
3print(model.weight.grad)
4
5loss2.backward()
6print(model.weight.grad)  # now includes both backward passes

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.

python
1accumulation_steps = 4
2optimizer.zero_grad()
3
4for step, (inputs, targets) in enumerate(loader):
5    outputs = model(inputs)
6    loss = criterion(outputs, targets) / accumulation_steps
7    loss.backward()
8
9    if (step + 1) % accumulation_steps == 0:
10        optimizer.step()
11        optimizer.zero_grad()

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() after backward() 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 .grad buffers 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.

Course illustration
Course illustration

All Rights Reserved.