Not change NaN if reset_if_nan is False

Author: Qianruipku

ferminet/train.py

@@ -979,7 +979,11 @@ def learning_rate_schedule(t_: jnp.ndarray) -> jnp.ndarray:
       loss = loss[0]
       # per batch variance isn't informative. Use weighted mean and variance
       # instead.
-      if not jnp.isnan(loss):
+      if cfg.optim.reset_if_nan and jnp.isnan(loss):
+        fallback_loss = weighted_stats.mean if weighted_stats is not None else 0.0
+        weighted_stats = statistics.exponentialy_weighted_stats(
+            alpha=0.1, observation=fallback_loss, previous_stats=weighted_stats)
+      else:
         weighted_stats = statistics.exponentialy_weighted_stats(
             alpha=0.1, observation=loss, previous_stats=weighted_stats)
       pmove = pmove[0]