How to debug NaN loss in PyTorch

A NaN loss in PyTorch means a training step produced a value that is not a finite number. It can come from invalid loss math, exploding gradients, bad input data, or mixed-precision overflow, and the first task is to stop the loop at the first non-finite tensor instead of letting later batches hide the source.

Autograd anomaly detection records the forward operation that later creates a bad backward value, while explicit torch.isfinite checks catch non-finite losses and gradients beside the training step. Enable both in a temporary debug run, then remove the extra checks after the failing operation is fixed because anomaly detection slows training.

A small CPU repro keeps the debugging pattern independent of private data and GPU state. The sample model applies sqrt() to negative predictions so the failing loss is deterministic; use the same guard positions in a copy of the broken training script before changing learning rates, optimizers, or model size.

1. Create a debug script with anomaly detection and finite-value checks.

   debug_nan_loss.py

   import argparse
    
   import torch
   from torch import nn
    
    
   def check_finite(name: str, tensor: torch.Tensor, *, epoch: int, batch: int) -> None:
       finite_mask = torch.isfinite(tensor)
       if bool(finite_mask.all()):
           return
    
       bad_values = tensor.detach()[~finite_mask].flatten()[:4].tolist()
       print(f"non_finite={name}")
       print(f"epoch={epoch} batch={batch} dtype={tensor.dtype} shape={tuple(tensor.shape)}")
       print(f"{name}_bad_values={bad_values}")
       print("hint=inspect the operation that created this tensor before changing the optimizer")
       raise SystemExit(1)
    
    
   def check_gradients(model: nn.Module, *, epoch: int, batch: int) -> None:
       for name, parameter in model.named_parameters():
           if parameter.grad is None:
               continue
           if not bool(torch.isfinite(parameter.grad).all()):
               bad_values = parameter.grad.detach()[~torch.isfinite(parameter.grad)].flatten()[:4].tolist()
               print(f"non_finite_gradient={name}")
               print(f"epoch={epoch} batch={batch} bad_values={bad_values}")
               raise SystemExit(1)
       print("gradient_check=finite")
    
    
   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 main() -> None:
       parser = argparse.ArgumentParser()
       parser.add_argument("--fixed", action="store_true")
       args = parser.parse_args()
    
       torch.manual_seed(7)
       torch.autograd.set_detect_anomaly(True, check_nan=True)
    
       features = torch.tensor(
           [
               [1.0, 0.5],
               [2.0, 1.0],
               [3.0, 1.5],
               [4.0, 2.0],
           ],
           dtype=torch.float32,
       )
       target = torch.tensor([[0.5], [1.0], [1.5], [2.0]], dtype=torch.float32)
    
       model = TinyRegressor()
       optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
    
       epoch = 1
       batch = 0
       optimizer.zero_grad(set_to_none=True)
       prediction = model(features)
       check_finite("prediction", prediction, epoch=epoch, batch=batch)
       print(f"prediction_min={prediction.min().item():.6f}")
       print(f"prediction_max={prediction.max().item():.6f}")
    
       if args.fixed:
           loss = nn.functional.mse_loss(prediction, target)
       else:
           loss = torch.sqrt(prediction).mean()
    
       check_finite("loss", loss, epoch=epoch, batch=batch)
       loss.backward()
       check_gradients(model, epoch=epoch, batch=batch)
       optimizer.step()
    
       print(f"loss={loss.item():.6f}")
       print("optimizer_step=complete")
    
    
   if __name__ == "__main__":
       main()

   torch.autograd.set_detect_anomaly(True, check_nan=True) makes autograd raise on NaN values created during backward. The explicit torch.isfinite checks stop earlier when the loss itself is already NaN or infinite.

2. Run the debug script against the failing loss path.

   $ python debug_nan_loss.py
   prediction_min=-1.850000
   prediction_max=-0.500000
   non_finite=loss
   epoch=1 batch=0 dtype=torch.float32 shape=()
   loss_bad_values=[nan]
   hint=inspect the operation that created this tensor before changing the optimizer

   The predictions are finite but negative, so the next invalid operation is the loss expression rather than the model output tensor.

3. Inspect the loss expression after the last finite tensor.

   The sample failure is torch.sqrt(prediction), which requires non-negative input. In a real training loop, check logarithms, square roots, divisions, normalization denominators, masking branches, and custom loss functions near the first non-finite report.

4. Apply a domain-safe loss fix in the debug copy.

   In the sample script, --fixed switches the loss from torch.sqrt(prediction) to nn.functional.mse_loss(prediction, target). In a real model, fix the invalid math or input contract instead of masking NaN after it appears.

5. Rerun the fixed training step.

   $ python debug_nan_loss.py --fixed
   prediction_min=-1.850000
   prediction_max=-0.500000
   gradient_check=finite
   loss=7.008750
   optimizer_step=complete

6. Remove the temporary debug script after the failing project has the guard pattern or the fix.

   $ rm debug_nan_loss.py