small fix

stanford-crfm · Apr 29, 2024 · 173454d · 173454d
1 parent e74cd02
commit 173454d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/optim/schedulefree_sophia.py b/src/levanter/optim/schedulefree_sophia.py
@@ -256,7 +256,7 @@ def _do_update():
             # EMAs of hessian
             nu = update_moment(new_hess, state.h, b2, 1)
             return ScaleBySophiaState(
-                count=state.count, hessian_count=state.hessian_count + 1, mu=state.mu, h=nu, hess_key=next_key
+                count=state.count, hessian_count=state.hessian_count + 1, z=state.z, h=nu, hess_key=next_key
             )
 
         def _dont_update():