diff --git a/single_shot_inference.py b/single_shot_inference.py index 79f4cb8d..a9392184 100644 --- a/single_shot_inference.py +++ b/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()