EMA params bug fix

rystrauss · Mar 9, 2022 · 3cef653 · 3cef653
1 parent 818e187
commit 3cef653
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bax/trainer.py b/bax/trainer.py
@@ -226,7 +226,7 @@ def _grad_step(
                 new_params,
             )
             new_ema_params = jmp.select_tree(
-                should_skip, new_ema_params, train_state.ema_params
+                should_skip, train_state.ema_params, new_ema_params
             )
         else:
             new_ema_params = train_state.ema_params