How to check tensors for NaN and Inf in TensorFlow

Invalid floating-point values often appear after a bad loss scale, division by zero, logarithm boundary, or exploding gradient in a TensorFlow run. Checking tensors for NaN and Inf stops the failure near the operation that produced the bad value instead of letting it spread into later losses, gradients, or exported outputs.

The targeted API is tf.debugging.check_numerics(), which checks one floating tensor and returns it unchanged when all values are finite. The broader API is tf.debugging.enable_check_numerics(), which makes eager execution and tf.function graph execution raise when an operation output contains NaN or Inf.

Numerics checks add debugging overhead and can produce long stack traces, so keep the global check around the failing run or a small reproduced batch. Disable the global mechanism after the focused check when the same Python process will continue with normal training or inference.

Steps to check TensorFlow tensors for NaN and Inf:

Open a terminal in the Python environment where TensorFlow imports successfully.
```
$ python3 -c "import tensorflow as tf; print(tf.__version__)"
2.21.0
```
Related: How to create a virtual environment for TensorFlow
Related: How to install TensorFlow with pip

Save the numerics check demo as

check_tensor_numerics.py

check_tensor_numerics.py

import os
 
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
import tensorflow as tf
 
tf.get_logger().setLevel("ERROR")
 
 
def print_numerics_error(label, error):
    for line in str(error).splitlines():
        cleaned = line.strip()
        wrapper_end = "}" * 2 + " "
        if cleaned.startswith("{" * 2 + "function_node") and wrapper_end in cleaned:
            cleaned = cleaned.split(wrapper_end, 1)[1]
        if "Detected Infinity or NaN" in cleaned or "Tensor had" in cleaned:
            cleaned = cleaned.replace("!!! ", "").replace(" !!!", "")
            cleaned = cleaned.replace(" (# of outputs: 1)", "")
            cleaned = cleaned.split(" [Op:", 1)[0]
            print(f"{label}: {cleaned}")
            return
    print(f"{label}: {error.__class__.__name__}")
 
 
print(f"TensorFlow {tf.__version__}")
 
safe_tensor = tf.debugging.check_numerics(
    tf.constant([1.0, 2.0], dtype=tf.float32),
    "input batch",
)
print(f"Safe tensor: {safe_tensor.numpy().tolist()}")
 
try:
    tf.debugging.check_numerics(
        tf.constant([1.0, float("inf")], dtype=tf.float32),
        "manual tensor check",
    )
except tf.errors.InvalidArgumentError as error:
    print_numerics_error("Manual check caught", error)
 
tf.debugging.enable_check_numerics()
 
try:
    tf.math.sqrt(tf.constant([4.0, -1.0], dtype=tf.float32))
except tf.errors.InvalidArgumentError as error:
    print_numerics_error("Eager check caught", error)
 
 
@tf.function
def unstable_log(values):
    return tf.math.log(values)
 
 
try:
    unstable_log(tf.constant([1.0, 0.0], dtype=tf.float32))
except tf.errors.InvalidArgumentError as error:
    print_numerics_error("tf.function check caught", error)
 
tf.debugging.disable_check_numerics()

tf.debugging.check_numerics() checks floating tensors such as float32 and float64. The message string should name the batch, loss, gradient, or layer so the error points back to the suspect part of the program.

Run the demo and confirm that the manual, eager, and tf.function checks all catch bad values.

$ python3 check_tensor_numerics.py
TensorFlow 2.21.0
Safe tensor: [1.0, 2.0]
Manual check caught: manual tensor check : Tensor had Inf values
Eager check caught: Detected Infinity or NaN in output 0 of eagerly-executing op "Sqrt"
tf.function check caught: Detected Infinity or NaN in output 0 of graph op "Log"

The safe tensor line shows that finite values pass through unchanged. The caught lines show an explicit Inf tensor, an eager Sqrt output that became NaN, and a tf.function Log output that became Inf.

Add a targeted check around the tensor that first becomes suspicious in the real program.
```
loss = loss_fn(labels, predictions)
tf.debugging.check_numerics(loss, "loss after forward pass")
```
Place targeted checks after numerically sensitive operations such as division, logarithms, exponentials, loss scaling, reduction, or gradient calculation.
Enable global numerics checking around a reproduced failing run when the source tensor is still unclear.
```
tf.debugging.enable_check_numerics()
 
try:
    for features, labels in train_dataset:
        train_step(features, labels)
finally:
    tf.debugging.disable_check_numerics()
```
enable_check_numerics() affects the current thread and raises tf.errors.InvalidArgumentError when eager or graph execution produces NaN or Inf.
Remove the temporary numerics demo after copying the check pattern into the failing project.
```
$ rm check_tensor_numerics.py
```