Batch 47,000 of a ResNet variant I was training for a medical imaging project. The loss had been dropping nicely: 1.4, then 1.1, then 0.87. Then nan. No error, no crash, just the computation graph quietly dying. Anyone who has spent time in deep learning knows how much of a day PyTorch NaN detection can swallow.
I did the obvious thing and flipped on torch.autograd.set_detect_anomaly(True). Training slowed to a crawl, roughly 10x on CPU and close to 100x on GPU. Three hours later it handed me a stack trace pointing at a layer that was perfectly healthy. The actual cause was a learning rate spike three layers upstream. The built-in tooling finds symptoms, not sources.
Why standard PyTorch NaN detection fails at scale
set_detect_anomaly is a blunt instrument. It puts PyTorch into synchronous mode and saves intermediate activations for every single operation so it can trace the gradient. It also tends to report where the NaN propagated to on the backward pass instead of where it originated on the forward pass.
That overhead does not survive contact with a production model, so I built a small detector on forward hooks. It costs about 3ms per pass, against the much larger latency the built-in anomaly engine adds.
The forward hook approach
register_forward_hook lets you inspect tensors as they flow through the network. Rather than replaying the whole graph, you check the output of each layer on the way past. torch.isnan().any() is a single CUDA kernel launch and finishes in microseconds.
def bbioon_nan_hook(module, inputs, output):
if torch.isnan(output).any() or torch.isinf(output).any():
# Capture the state immediately
stats = {
"layer": module.__class__.__name__,
"mean": output[~torch.isnan(output)].mean().item(),
"batch_idx": getattr(module, "_current_batch", -1)
}
raise ValueError(f"NaN detected at {stats['layer']} in batch {stats['batch_idx']}")Making it safe in a real training loop
With multi-worker DataLoaders or a complicated training loop, a print statement will not hold up. You need thread safety and some care with memory. Record events across threads without a lock and the race conditions will be worse to debug than the original bug. Part of how good debugging skills can make you a better developer is knowing which boring parts to automate so your attention stays on the architecture.
My version uses a threading.Lock() so the first occurrence gets logged cleanly without another thread stepping on it. There is also a gradient norm guard. Exploding gradients come before most NaNs, so checking param.grad.norm() after loss.backward() often catches the explosion a full step before it corrupts your weights.
What the overhead actually is
On a 5-layer MLP, the forward hook method measured around 3 to 4 ms of overhead in my benchmarks. set_detect_anomaly came in at 7 to 8 ms on CPU. On a GPU-accelerated Transformer the gap stops being a matter of milliseconds and becomes the difference between training finishing today and training finishing next week.
If PyTorch NaN detection is eating your dev hours, hand it to me. I have been wrestling with WordPress since the 4.x days, and I bring the same pragmatism to every stack I touch.
What to put in your training pipeline
Leave the nuclear option alone. When training gets unstable, register a hook, watch your gradient norms, and build sequential models with OrderedDict so layer names mean something in the logs. What you want is a setup that tells you where the error happened before it spreads through the rest of the graph. The official autograd mechanics cover the rest.
