🎓 A1Learner
🔍

PyTorch Hunting a NaN Loss

Hunting a NaN loss means finding the first tensor that stops being finite. The hunt checks tensors in the order they are made, instead of trusting the final broken loss.

A rate of 50 breaks the run

The pain comes first. A learning rate of 50 makes each step overshoot further than the last.

torch.manual_seed(0)
x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[3.0], [5.0], [7.0], [9.0]])
model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=50.0)
for step in range(3):
    loss = loss_fn(model(x), y)
    print(f"Step {step} loss: {loss.item():.4g}")
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Check a tensor for bad values

isfinite tests every entry at once. One infinite entry makes the all-finite check fail.

values = torch.tensor([1.5, float("inf"), 0.25])
print("Each finite:", torch.isfinite(values).tolist())
print("All finite:", bool(torch.isfinite(values).all()))

A hook checks every pass

A forward hook sees each output as it is made. The hook stores one flag, and the training loop reads the flag after the pass.

torch.manual_seed(0)
model = nn.Linear(3, 2)
seen = {}
def check(module, inputs, output):
    seen["finite"] = bool(torch.isfinite(output).all())
handle = model.register_forward_hook(check)
out = model(torch.randn(4, 3))
print("Hook saw output finite:", seen["finite"])
handle.remove()

Example

import torch
from torch import nn

torch.manual_seed(0)
x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[3.0], [5.0], [7.0], [9.0]])

model = nn.Linear(1, 1)
seen = {}


def check_hook(module, inputs, output):
    seen["prediction_finite"] = bool(torch.isfinite(output).all())


handle = model.register_forward_hook(check_hook)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=50.0)

first_bad = None
for step in range(20):
    pred = model(x)
    loss = loss_fn(pred, y)
    print(f"Step {step} loss: {loss.item():.4g}")
    if not seen["prediction_finite"]:
        first_bad = (step, "prediction")
        break
    if not torch.isfinite(loss):
        first_bad = (step, "loss")
        break
    optimizer.zero_grad()
    loss.backward()
    if not torch.isfinite(model.weight.grad).all():
        first_bad = (step, "weight gradient")
        break
    optimizer.step()

handle.remove()
print("First bad step:", first_bad[0])
print("First bad tensor:", first_bad[1])
print("Prediction finite at that step:", seen["prediction_finite"])

Key points

  • NaN: a value that means not a number, reached after training numbers overflow
  • Finite: an ordinary number, not infinity and not NaN
  • Use torch.isfinite(t) to test which entries of a tensor are still finite
  • A huge learning rate can grow the weights until the squared loss overflows
  • The first bad tensor is the earliest check that fails, not the last value printed
  • Checking prediction, then loss, then gradient follows the order a step makes them

In this example, the hunt runs the broken training from the first snippet, with checks added. A forward hook records whether each prediction is fully finite. Every step prints its loss, then the loop tests the prediction flag, the loss, and the weight gradient in order. The losses grow by about six orders of magnitude per step. At step 7 the loss prints inf, and the loop stops there. The prediction at step 7 is still finite, as the last line shows. Squaring those huge predictions inside the mean squared error is what overflowed. So the first bad tensor is the loss at step 7, and no later NaN needs explaining.

Note

In this run the loss reaches infinity before any NaN, so the hunt stops at the first non-finite value

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A run prints a NaN loss at step 40. What should the hunt find?