check_nan fn corrected

Author: prishajain1

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -41,11 +41,20 @@
 BlockSizes = common_types.BlockSizes
 
 def check_nan(tensor: jax.Array, name: str):
+    if tensor is None:
+        # jax.debug.print works fine with regular python strings and values
+        print(f"[DEBUG NaN Check] {name} on process {jax.process_index()}: Tensor is None")
+        return
+
     has_nans = jnp.isnan(tensor).any()
     has_infs = jnp.isinf(tensor).any()
-    # Use jax.debug.print to print during JITted execution
+
+    # Pass the JAX arrays (has_nans, has_infs) as kwargs
+    # Use placeholders {} in the f-string for these runtime values
     jax.debug.print(f"[DEBUG NaN Check] {name} on process {jax.process_index()}: "
-                    f"Has NaNs: {has_nans}, Has Infs: {has_infs}")
+                    "Has NaNs: {has_nans_val}, Has Infs: {has_infs_val}",
+                    has_nans_val=has_nans, has_infs_val=has_infs)
+
 
 def get_frequencies(max_seq_len: int, theta: int, attention_head_dim: int, use_real: bool):
   h_dim = w_dim = 2 * (attention_head_dim // 6)