Add per-layer NaN check to track where values diverge

single_shot_inference.py

@@ -391,6 +391,13 @@ def main():
                 print(f"  L{li}: no weights!")
                 continue
             x = forward_layer(x, lw, li, cfg, rope_cos, rope_sin)
+            if li % 5 == 0 or torch.isnan(x.float()).any():
+                has_nan = torch.isnan(x.float()).any().item()
+                xmax = x.float().abs().max().item()
+                print(f"    L{li}: nan={has_nan}, max_abs={xmax:.4f}")
+                if has_nan and xmax == 0:
+                    print(f"    L{li}: all NaN, stopping")
+                    break
             del lw
             if li % 10 == 9:
                 reader.clear()