small fix

stanford-crfm · Dec 19, 2024 · c47c4c5 · c47c4c5
1 parent 5bafdcf
commit c47c4c5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/optim/kron.py b/src/levanter/optim/kron.py
@@ -497,7 +497,7 @@ def _balance_Q(Q: List[jax.Array]):
         updates = grads_structure.unflatten(precond_gs)
         Qs = grads_structure.unflatten(Qs)
 
-        precond_gs = updates_struct.unflatten(precond_gs)
+        updates = updates_struct.unflatten(updates)
 
         # dtypes and new state
         mu = otu.tree_cast(mu, mu_dtype)