Merge pull request #104 from Qianruipku:hotfix

jsspencer · jsspencer · commit 6848678dc716 · 2025-09-19T10:38:05.000+01:00
PiperOrigin-RevId: 808807612
Change-Id: Iab503acc70b075e87e2a1f3fd4ff452547db265b
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`