fix: prevent NaN loss from updating weighted statistics

Author: Qianruipku

ferminet/train.py

@@ -979,8 +979,9 @@ def learning_rate_schedule(t_: jnp.ndarray) -> jnp.ndarray:
       loss = loss[0]
       # per batch variance isn't informative. Use weighted mean and variance
       # instead.
-      weighted_stats = statistics.exponentialy_weighted_stats(
-          alpha=0.1, observation=loss, previous_stats=weighted_stats)
+      if not jnp.isnan(loss):
+        weighted_stats = statistics.exponentialy_weighted_stats(
+            alpha=0.1, observation=loss, previous_stats=weighted_stats)
       pmove = pmove[0]
 
       # Update observables