Introduction

If a TensorFlow model predicts normally before saving but produces NaN only after restore, the saved weights are not the only thing to inspect. The failure often comes from a mismatch in how variables, optimizer state, preprocessing, or training-mode behavior are reconstructed. The fastest way to debug it is to compare the pre-save and post-restore pipeline step by step instead of guessing about “corrupt checkpoints”.

First Verify That the Restored Model Is Actually Equivalent

Start by checking whether the restored model produces the same finite output on the same input. Do this with a tiny controlled example before involving your full dataset.

1import tensorflow as tf
2import numpy as np
3
4model = tf.keras.Sequential([
5    tf.keras.layers.Dense(8, activation="relu"),
6    tf.keras.layers.Dense(1)
7])
8
9x = np.array([[1.0, 2.0]], dtype=np.float32)
10_ = model(x)
11
12before = model(x)
13model.save_weights("demo.weights.h5")
14
15restored = tf.keras.Sequential([
16    tf.keras.layers.Dense(8, activation="relu"),
17    tf.keras.layers.Dense(1)
18])
19_ = restored(x)
20restored.load_weights("demo.weights.h5")
21after = restored(x)
22
23print(before.numpy())
24print(after.numpy())

If this small comparison already breaks, the problem is in model construction or checkpoint compatibility. If it works, the issue is likely elsewhere in the full inference path.

Build the Variables Before Loading When Using Custom Models

Subclassed models and custom layers can behave differently from plain Sequential models because variables may not exist until the first call. If the restore path does not match the build path, some values may remain unset or deferred in surprising ways.

A reliable habit is to create the variables before loading weights.

restored = MyModel()
_ = restored(tf.zeros((1, 16)))
restored.load_weights("model.weights.h5")

That makes the variable structure explicit. It also reveals mismatches sooner, because TensorFlow can compare the saved tensors against an already-built model.

Check Training Versus Inference Behavior

NaN after restore is often caused by using the model in a different mode. Layers such as batch normalization and dropout behave differently during training and inference.

y = restored(x, training=False)
tf.debugging.assert_all_finite(y, "non-finite output after restore")

If the original path used training=False but the restored path accidentally uses training=True, moving statistics, randomness, or unstable batch behavior can create output differences that look like checkpoint corruption.

This matters especially when the model is wrapped inside a custom prediction function that does more than a direct forward pass.

Inspect the Input Pipeline, Not Only the Model

Many restore-only NaN bugs are actually input bugs. A different preprocessing step, a missing normalization constant, or a changed dtype can push values into invalid ranges after deployment.

Use finite checks on the input as well as the output.

1features = tf.convert_to_tensor([[1.0, 2.0]], dtype=tf.float32)
2tf.debugging.assert_all_finite(features, "non-finite input")
3preds = restored(features, training=False)
4tf.debugging.assert_all_finite(preds, "non-finite prediction")

If the restored code path loads scalers, vocabularies, or lookup tables separately, verify those artifacts too. A model can be restored perfectly and still output NaN because one preprocessing asset is missing or inconsistent.

Optimizer State Matters Only If You Resume Training

If NaN appears when resuming training rather than during plain inference, check whether optimizer slots were restored. Restoring only model weights but not optimizer state can change the next update step dramatically, especially with adaptive optimizers.

For inference-only debugging, ignore the optimizer and focus on model variables plus input consistency. For training continuation, restore the full checkpoint object that includes optimizer state.

Compare Weights Directly When Necessary

If you still suspect the checkpoint, compare the saved and restored weights numerically.

for a, b in zip(model.get_weights(), restored.get_weights()):
    print(np.allclose(a, b))

If all arrays match, the problem is almost certainly outside raw weight serialization. That is a useful boundary because it turns a vague debugging problem into a concrete one.

Common Pitfalls

• Loading weights into a model that has not been built the same way as the original model.
• Accidentally running inference with training=True.
• Restoring the model but not the preprocessing assets or normalization rules.
• Assuming any restore-time NaN must mean the checkpoint file is corrupted.
• Resuming training without restoring optimizer state.

Summary

• Restore-only NaN issues are usually caused by reconstruction mismatches, not magic checkpoint failure.
• Compare one fixed input before save and after restore to localize the problem.
• Build subclassed models before loading weights.
• Check inference mode, preprocessing, and asset loading as carefully as the weights themselves.
• If weights match numerically, move your debugging effort to the surrounding pipeline.