Repro: SRGAN Generator block8_final_conv NaN on GPU vs finite on CPU (minimal test bundle)

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

2.22.0-dev20260302

Custom code

Yes

OS platform and distribution

Linux Ubuntu 22.04 (Docker environment)

Mobile device

N/A

Python version

3.10

Bazel version

N/A

GCC/compiler version

N/A

CUDA/cuDNN version

12.5.1/9.2.1

GPU model and memory

NVIDIA GeForce RTX 3090, 24GB

Current behavior?

Current behavior

I reproduced the issue on the SRGAN Generator with a minimal, CIPIHunter-aligned setup: the same attack batch from iter_008.npz and the same Generator weights (bundled in srgan9_repro.npz), using one inlined Generator implementation. I run the full forward pass twice—once on GPU and once on CPU. On the output of block8_final_conv (the last Conv2D in the Generator, i.e. Generator:0.Conv2D:2 in flattened naming), the GPU run reports NaNs (has_nan: True), while the CPU run is finite with no NaNs (has_nan: False). I do not see Infs on either side in this check. Example stack: TensorFlow / Keras, cuDNN 9.x, RTX 3090 (compute capability 8.6).

Expected behavior

Under float32, with the same weights and the same input, I expect GPU and CPU to stay numerically consistent in the sense that we should not get NaNs only on GPU while the CPU result remains finite—unless there is documented, acceptable backend-specific behavior. Small floating-point drift is fine; NaN on one device only is not expected for this forward pass when the input and weights are finite.

Here are the external files required for reproducing this issue:

demo4.zip

Standalone code to reproduce the issue

import math
import numpy as np
import tensorflow as tf
tf.keras.mixed_precision.set_global_policy("float32")


class PixelShuffle(tf.keras.layers.Layer):
    def __init__(self, s):
        super().__init__()
        self.s = s
    def call(self, x):
        return tf.nn.depth_to_space(x, self.s)

class ResidualBlock(tf.keras.layers.Layer):
    def __init__(self, ch):
        super().__init__()
        self.pad1 = tf.keras.layers.ZeroPadding2D(1)
        self.conv1 = tf.keras.layers.Conv2D(ch, 3, padding="valid", data_format="channels_last")
        self.bn1 = tf.keras.layers.BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5)
        self.prelu = tf.keras.layers.PReLU(shared_axes=[1, 2])
        self.pad2 = tf.keras.layers.ZeroPadding2D(1)
        self.conv2 = tf.keras.layers.Conv2D(ch, 3, padding="valid", data_format="channels_last")
        self.bn2 = tf.keras.layers.BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5)
        self.add = tf.keras.layers.Add()
    def call(self, x, training=False):
        r = self.conv1(self.pad1(x))
        r = self.prelu(self.bn1(r, training=training))
        r = self.bn2(self.conv2(self.pad2(r)), training=training)
        return self.add([x, r])

class UpsampleBLock(tf.keras.layers.Layer):
    def __init__(self, ch, s):
        super().__init__()
        self.pad = tf.keras.layers.ZeroPadding2D(1)
        self.conv = tf.keras.layers.Conv2D(ch * s**2, 3, padding="valid", data_format="channels_last")
        self.ps = PixelShuffle(s)
        self.prelu = tf.keras.layers.PReLU(shared_axes=[1, 2])
    def call(self, x, training=False):
        return self.prelu(self.ps(self.conv(self.pad(x))))

class Generator(tf.keras.Model):
    def __init__(self, sf):
        super().__init__()
        n_up = int(math.log(sf, 2))
        self.block1_pad = tf.keras.layers.ZeroPadding2D(4)
        self.block1_conv = tf.keras.layers.Conv2D(64, 9, padding="valid", data_format="channels_last")
        self.block1_prelu = tf.keras.layers.PReLU(shared_axes=[1, 2])
        self.block2 = ResidualBlock(64)
        self.block3 = ResidualBlock(64)
        self.block4 = ResidualBlock(64)
        self.block5 = ResidualBlock(64)
        self.block6 = ResidualBlock(64)
        self.block7_pad = tf.keras.layers.ZeroPadding2D(1)
        self.block7_conv = tf.keras.layers.Conv2D(64, 3, padding="valid", data_format="channels_last")
        self.block7_bn = tf.keras.layers.BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5)
        self.block7_add = tf.keras.layers.Add()
        self.block8_layers = [UpsampleBLock(64, 2) for _ in range(n_up)]
        self.block8_final_pad = tf.keras.layers.ZeroPadding2D(4)
        self.block8_final_conv = tf.keras.layers.Conv2D(3, 9, padding="valid", data_format="channels_last")
        self.captured = {}

    def call(self, x, training=False):
        b1 = self.block1_prelu(self.block1_conv(self.block1_pad(x)))
        b = self.block2(b1, training=training)
        b = self.block3(b, training=training)
        b = self.block4(b, training=training)
        b = self.block5(b, training=training)
        b = self.block6(b, training=training)
        b7 = self.block7_bn(self.block7_conv(self.block7_pad(b)), training=training)
        b8 = self.block7_add([b1, b7])
        for layer in self.block8_layers:
            b8 = layer(b8, training=training)
        b8 = self.block8_final_conv(self.block8_final_pad(b8))
        self.captured["conv2_out"] = tf.cast(b8, tf.float32).numpy()
        return (tf.math.tanh(b8) + 1.0) / 2.0


d = np.load("srgan9_repro.npz", allow_pickle=True)
x_input = d["input"]
gen_w = list(d["gen_w"])

def run(device):
    tf.keras.backend.clear_session()
    with tf.device(f"/{device.upper()}:0"):
        g = Generator(sf=4)
        g(tf.zeros((1, 64, 64, 3), tf.float32), training=False)
        g.set_weights(gen_w)
        g(tf.convert_to_tensor(x_input, tf.float32), training=False)
        return g.captured["conv2_out"]

out_g = run("gpu")
out_c = run("cpu")
print("GPU block8_final_conv max:", np.max(out_g), "has_nan:", bool(np.isnan(out_g).any()), "has_inf:", bool(np.isinf(out_g).any()))
print("CPU block8_final_conv max:", np.max(out_c), "has_nan:", bool(np.isnan(out_c).any()), "has_inf:", bool(np.isinf(out_c).any()))

Relevant log output

GPU block8_final_conv max: nan has_nan: True has_inf: False
CPU block8_final_conv max: 6.2258007e+37 has_nan: False has_inf: False