How to debug NaN loss in PyTorch

A training run can cross from valid arithmetic into NaN or infinity in a single operation, while later batches only repeat the damage. The useful boundary is the last finite tensor before the first non-finite loss or gradient, because that boundary narrows the fault to one part of the forward or backward pass.

The torch.isfinite function checks tensors directly, and autograd anomaly detection traces a failing backward operation to the forward call that created it. Anomaly detection adds runtime overhead, so it belongs in a short debug run rather than the normal training loop.

The CPU sample uses negative predictions with torch.sqrt to reproduce a non-finite loss without private data or GPU state. Its guards sit immediately after the model output, after the loss calculation, and after backward(); those are the same boundaries to keep when adapting the pattern to a failing batch.

1. Create /debug_nan_loss.py with a finite-tensor guard.

   debug_nan_loss.py

   import argparse
    
   import torch
   from torch import nn
    
    
   def check_finite(name: str, tensor: torch.Tensor, *, step: int) -> None:
       finite = torch.isfinite(tensor)
       if bool(finite.all()):
           return
    
       values = tensor.detach()[~finite].flatten()[:4].tolist()
       print(f"non_finite={name}")
       print(f"step={step} shape={tuple(tensor.shape)} dtype={tensor.dtype}")
       print(f"values={values}")
       raise SystemExit(1)

   The guard exits at the first NaN, positive infinity, or negative infinity instead of allowing the optimizer to update parameters from invalid values.

2. Add a gradient guard below check_finite.

   def check_gradients(model: nn.Module, *, step: int) -> None:
       for name, parameter in model.named_parameters():
           if parameter.grad is not None:
               check_finite(f"gradient:{name}", parameter.grad, step=step)
       print("gradients=finite")

3. Add the deterministic model fixture below check_gradients.

   class TinyRegressor(nn.Module):
       def __init__(self) -> None:
           super().__init__()
           self.linear = nn.Linear(2, 1)
           with torch.no_grad():
               self.linear.weight.copy_(torch.tensor([[-0.35, -0.20]]))
               self.linear.bias.fill_(-0.05)
    
       def forward(self, inputs: torch.Tensor) -> torch.Tensor:
           return self.linear(inputs)
    
    
   def build_fixture() -> tuple[nn.Module, torch.Tensor, torch.Tensor]:
       features = torch.tensor(
           [[1.0, 0.5], [2.0, 1.0], [3.0, 1.5], [4.0, 2.0]]
       )
       targets = torch.tensor([[0.5], [1.0], [1.5], [2.0]])
       return TinyRegressor(), features, targets

   The fixed weights make every prediction negative, which gives the invalid square-root branch a repeatable failure surface.

4. Add the guarded training step below build_fixture.

   def run_step(fixed: bool) -> None:
       torch.manual_seed(7)
       model, features, targets = build_fixture()
       optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
       optimizer.zero_grad(set_to_none=True)
       step = 0
    
       with torch.autograd.set_detect_anomaly(True, check_nan=True):
           prediction = model(features)
           check_finite("prediction", prediction, step=step)
           print(
               f"prediction_range={prediction.min().item():.6f},"
               f"{prediction.max().item():.6f}"
           )
    
           if fixed:
               loss = nn.functional.mse_loss(prediction, targets)
           else:
               loss = torch.sqrt(prediction).mean()
    
           check_finite("loss", loss, step=step)
           loss.backward()
    
       check_gradients(model, step=step)
       optimizer.step()
       print(f"loss={loss.item():.6f}")
       print("optimizer_step=complete")

   The output and loss checks cover forward arithmetic, while anomaly detection and the gradient guard cover failures produced during backward().

5. Add the command-line entry point below run_step.

   parser = argparse.ArgumentParser()
   parser.add_argument("--fixed", action="store_true")
   args = parser.parse_args()
   run_step(args.fixed)

6. Run the failing loss path to stop at the first non-finite tensor.

   $ python debug_nan_loss.py
   prediction_range=-1.850000,-0.500000
   non_finite=loss
   step=0 shape=() dtype=torch.float32
   values=[nan]

7. Compare the last finite tensor with the first non-finite report.

   The finite negative predictions followed by a non-finite loss identify torch.sqrt(prediction) as the invalid operation. The --fixed branch replaces that expression with nn.functional.mse_loss; a project failure may instead point to a logarithm, division, normalization denominator, mask, input batch, or mixed-precision range.

8. Run the corrected path to confirm finite gradients and a completed optimizer step.

   $ python debug_nan_loss.py --fixed
   prediction_range=-1.850000,-0.500000
   gradients=finite
   loss=7.008750
   optimizer_step=complete

9. Remove the temporary debug script after the project training step passes the same finite-value checks.

   $ rm debug_nan_loss.py