Move eps outside sqrt

hr0nix · Sep 18, 2022 · 9e9658a · 9e9658a
1 parent 41ecc29
commit 9e9658a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optax_adan/transform.py b/optax_adan/transform.py
@@ -73,7 +73,7 @@ def update_fn(updates, state, params=None):
         n_hat = bias_correction(n, decay_n, count_inc)
 
         new_updates = jax.tree_util.tree_map(
-            lambda mm, vv, nn: (mm + decay_v * vv) / jnp.sqrt(nn + eps), m_hat, v_hat, n_hat)
+            lambda mm, vv, nn: (mm + decay_v * vv) / (jnp.sqrt(nn) + eps), m_hat, v_hat, n_hat)
 
         return new_updates, ScaleByAdanState(count=count_inc, m=m, v=v, n=n, prev_grad=updates)