How to zero gradients in PyTorch

Parameter gradients in a PyTorch model remain on each trainable parameter after loss.backward() runs. The next backward pass adds into those buffers unless they are cleared, so ordinary mini-batch training needs a reset before each new optimizer update.

optimizer.zero_grad() resets the parameters managed by that optimizer. In current PyTorch, the default set_to_none=True stores None instead of zero-filled gradient tensors, which reduces memory work and lets the next backward pass create fresh gradient tensors.

Call the reset before loss.backward() for a normal one-batch-one-step loop. Delay it only when intentionally accumulating gradients across microbatches, and clip gradients after backward() but before optimizer.step() when clipping is part of the training loop.

1. Create a minimal gradient-reset smoke script.

   gradient_zero_demo.py

   import torch
   from torch import nn
    
    
   torch.manual_seed(11)
    
   features = torch.tensor(
       [
           [0.0, 0.5],
           [1.0, -0.5],
           [2.0, 1.0],
       ],
       dtype=torch.float32,
   )
   targets = torch.tensor([[0.2], [0.7], [1.6]], dtype=torch.float32)
    
   model = nn.Linear(2, 1)
   loss_fn = nn.MSELoss()
   optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
    
    
   def grad_summary(label):
       pieces = []
       for name, parameter in model.named_parameters():
           if parameter.grad is None:
               pieces.append(f"{name}=None")
           else:
               pieces.append(
                   f"{name}_norm={parameter.grad.norm().item():.4f}; "
                   f"{name}_sum={parameter.grad.sum().item():.4f}"
               )
       print(f"{label}: {' | '.join(pieces)}")
    
    
   grad_summary("before backward")
    
   loss = loss_fn(model(features), targets)
   loss.backward()
   grad_summary("after backward")
    
   optimizer.step()
   optimizer.zero_grad()
   grad_summary("after zero_grad default")
    
   loss = loss_fn(model(features), targets)
   loss.backward()
   grad_summary("after next backward")
    
   optimizer.zero_grad(set_to_none=False)
   grad_summary("after zero_grad set_to_none_false")

   The script keeps the model small enough to inspect gradient buffers directly. optimizer.zero_grad() uses the current default set_to_none=True, and the final call shows the zero-tensor behavior requested with set_to_none=False.

2. Run the smoke script.

   $ python gradient_zero_demo.py
   before backward: weight=None | bias=None
   after backward: weight_norm=3.0538; weight_sum=-3.8573 | bias_norm=1.2983; bias_sum=-1.2983
   after zero_grad default: weight=None | bias=None
   after next backward: weight_norm=2.3521; weight_sum=-2.9601 | bias_norm=0.8466; bias_sum=-0.8466
   after zero_grad set_to_none_false: weight_norm=0.0000; weight_sum=0.0000 | bias_norm=0.0000; bias_sum=0.0000

3. Check the gradient state transitions in the output.

   before backward shows no stored gradients. after backward shows populated gradient tensors. after zero_grad default shows None again, and after zero_grad set_to_none_false shows explicit zero tensors.

4. Call optimizer.zero_grad() before loss.backward() in each normal training iteration.

   for batch_features, batch_targets in train_loader:
       optimizer.zero_grad()
       predictions = model(batch_features)
       loss = loss_fn(predictions, batch_targets)
       loss.backward()
       optimizer.step()

   Use optimizer.zero_grad() for the parameters that optimizer will step. model.zero_grad() also clears module parameter gradients, but it can drift from the optimizer when a loop uses selected parameter groups.

5. Keep the default reset mode unless later code needs gradient tensors.

   optimizer.zero_grad()
    
   # Use this only when code expects parameter.grad to be a zero tensor.
   optimizer.zero_grad(set_to_none=False)

   With set_to_none=True, parameters that receive no gradient stay at None after the next backward pass, and torch.optim skips optimizer work for those parameters. With zero tensors, the optimizer sees a gradient value of zero.

6. Remove the temporary smoke script.

   $ rm gradient_zero_demo.py