Fix multi-device NaN check in KFAC training step

Qianruipku · Qianruipku · commit 48931c6892e9 · 2025-08-08T17:06:57.000+01:00
Replace direct boolean check with jnp.any() to handle loss arrays
from multiple devices when detecting NaN values.
diff --git a/ferminet/train.py b/ferminet/train.py
@@ -378,7 +378,7 @@ def step(
         damping=shared_damping,
     )
 
-    if reset_if_nan and jnp.isnan(stats['loss']):
+    if reset_if_nan and jnp.any(jnp.isnan(stats['loss'])):
       new_params = old_params
       new_state = old_state
     return data, new_params, new_state, stats['loss'], stats['aux'], pmove

Original file line number	Diff line number	Diff line change
`@@ -378,7 +378,7 @@ def step(`
`378`	`378`	`damping=shared_damping,`
`379`	`379`	`)`
`380`	`380`
`381`		`- if reset_if_nan and jnp.isnan(stats['loss']):`
	`381`	`+ if reset_if_nan and jnp.any(jnp.isnan(stats['loss'])):`
`382`	`382`	`new_params = old_params`
`383`	`383`	`new_state = old_state`
`384`	`384`	`return data, new_params, new_state, stats['loss'], stats['aux'], pmove`