In a PyTorch training loop, why do you need to call optimizer.zero_grad() before backpropagation?

How to think about it

For the usual one-batch, one-update loop, call optimizer.zero_grad() before loss.backward() so the update uses only the current batch’s gradient. PyTorch adds new gradients to each parameter’s existing .grad buffer rather than overwriting it; the exception is when you deliberately accumulate several batches before one optimizer step.

Why PyTorch accumulates gradients

PyTorch’s autograd system, which records tensor operations and differentiates them during backpropagation, stores the gradient of each trainable parameter in that parameter’s .grad attribute.

When you run:

loss.backward()

PyTorch computes how much each parameter contributed to the loss and writes those derivatives into .grad. If .grad already contains a tensor, PyTorch adds the new derivative to it.

That behaviour is intentional. It lets you build a gradient from several backward passes. For example, you may need to:

• accumulate gradients over multiple small batches because a large batch does not fit in GPU memory;
• compute separate loss terms in separate backward passes;
• delay an optimizer update until several pieces of work are complete.

The training loop therefore has to say when one gradient calculation ends and the next one begins. optimizer.zero_grad() marks that boundary by clearing the old gradient buffers.

An optimizer does not normally inspect the entire computation graph. It reads the gradients currently stored in the parameters and uses them to update the weights. For plain stochastic gradient descent, the rule is:

new_weight = old_weight - learning_rate * stored_gradient

If the stored gradient contains leftovers, the optimizer cannot tell which part came from the current batch.

A numerical example

Suppose the model has one parameter, w, and the learning rate is 0.1. Start with w = 0.

The first batch produces a gradient of +3.

With a clean gradient buffer:

w = 0 - 0.1 × 3 = -0.3

Now the second batch produces a gradient of -1. The correct update should use only that second gradient:

w = -0.3 - 0.1 × (-1) = -0.2

The second batch wants to move w upward.

If you forget to clear the gradient, the second backward pass adds -1 to the old +3:

stored gradient = +3 + (-1) = +2

The optimizer then performs this update:

w = -0.3 - 0.1 × 2 = -0.5

The model moves in the wrong direction because the previous batch is still voting. In a real network, this happens independently for millions or billions of parameter values.

The ordinary training-loop order

A standard loop looks like this:

for inputs, targets in train_loader:
    optimizer.zero_grad()

    predictions = model(inputs)
    loss = criterion(predictions, targets)

    loss.backward()
    optimizer.step()

The important boundary is between optimizer steps:

1. Clear gradients left by the previous update.
2. Run the forward pass and calculate the current loss.
3. Run backpropagation to fill .grad with the current batch’s signal.
4. Let the optimizer read those gradients and update the parameters.

The clearing call can technically happen immediately before loss.backward() rather than before the forward pass. The forward pass does not change parameter gradients. Putting it at the top of the loop is conventional because it makes the loop’s state explicit and leaves less room for an accidental second backward pass.

Do not put zero_grad() between backward() and step(). That erases the gradients before the optimizer can use them:

loss.backward()
optimizer.zero_grad()  # The update now has nothing useful to read.
optimizer.step()

The model may appear to train, but its parameters will usually stop changing because the optimizer sees cleared gradients.

What zero_grad() actually clears

optimizer.zero_grad() clears the gradients for the parameters registered with that optimizer. It does not:

• reset the model’s weights;
• undo the previous optimizer update;
• reset Adam’s moving averages;
• reset momentum;
• remove the computation graph from the forward pass.

It only clears the current derivative buffers.

A cleared gradient can be represented either as a tensor full of zeros or as None. PyTorch supports:

optimizer.zero_grad(set_to_none=True)

With set_to_none=True, a parameter’s .grad becomes None until a later backward pass gives it a gradient. This can avoid writing a large tensor of zeros and lets an optimizer distinguish between “this parameter received a zero gradient” and “this parameter received no gradient at all.” The exact update behaviour for that distinction depends on the optimizer, so code should not treat None and a zero tensor as universally interchangeable.

optimizer.zero_grad() and model.zero_grad() are often interchangeable in a simple model with one optimizer. The distinction matters when an optimizer contains only some model parameters, or when several optimizers share parameters. In those cases, clear the parameters that participate in the update you are about to perform.

The important exception: intentional gradient accumulation

You should not clear gradients before every backward pass if you intentionally want several batches to contribute to one update.

This is common when the GPU can process only four examples at a time, but you want an effective batch size of sixteen. You can process four microbatches, accumulate their gradients, and call optimizer.step() once:

accumulation_steps = 4
optimizer.zero_grad(set_to_none=True)

for step, (inputs, targets) in enumerate(train_loader):
    predictions = model(inputs)
    loss = criterion(predictions, targets)

    (loss / accumulation_steps).backward()

    if (step + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

The division matters when criterion returns the mean loss for each microbatch. Without it, four mean gradients would be added together, producing roughly four times the gradient of their average. If the final accumulation window has fewer than four microbatches, production code must handle that window deliberately; blindly dividing by four makes the final update smaller than intended.

Gradient accumulation approximates a larger batch. It is not perfectly identical to using a physically larger batch. Batch-normalisation statistics, dropout masks, optimizer state updates, and the timing of gradient clipping can differ. That distinction matters when reproducing experiments or tuning a training run.

The failure mode is easy to create: put zero_grad() inside the microbatch loop. Then every microbatch erases the previous one, and only the last microbatch affects the update. You intended an effective batch of sixteen examples; you silently trained on four.

A subtle point about calling it after step()

This ordering also works in a basic loop:

for inputs, targets in train_loader:
    predictions = model(inputs)
    loss = criterion(predictions, targets)

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

After step(), the gradients are no longer needed, so clearing them there is valid as long as the next backward pass cannot happen first. Most teams put zero_grad() before the next forward or backward pass because that makes the invariant obvious: every ordinary update starts with an empty gradient buffer.

The mistake is not specifically “zeroing after the step.” The mistake is allowing an unwanted old gradient to survive into the next backward pass.

What they’ll ask next

Can I call model.zero_grad() instead?
Usually, yes. optimizer.zero_grad() clears parameters known to that optimizer, while model.zero_grad() clears parameters in the module. With multiple optimizers or shared parameters, choose deliberately rather than assuming they clear the same set.

Does zero_grad() reset momentum or Adam’s state?
No. It clears only the current gradients. Momentum buffers and Adam’s first- and second-moment estimates remain, because those are part of the optimizer’s history and are updated separately.

How do I accumulate gradients across microbatches correctly?
Do not clear between microbatches. Scale each microbatch loss so the accumulated gradient has the intended magnitude, call backward() for each one, call optimizer.step() once per accumulation window, and then clear the gradients before the next window.

Say this in the interview

“PyTorch accumulates derivatives in each parameter’s .grad, so in a normal one-batch-per-update loop I call optimizer.zero_grad() before loss.backward() to prevent stale gradients from previous batches changing the current update; I omit that clearing step only when I am deliberately accumulating gradients.”