whoops

stanford-crfm · Dec 15, 2024 · f7f2382 · f7f2382
1 parent 3a06e1c
commit f7f2382
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/optim/kron.py b/src/levanter/optim/kron.py
@@ -896,7 +896,7 @@ def pass_through_fn(key, qs, grads_in, bal_counter):
             do_update,
             update_preconditioner_fn,
             pass_through_fn,
-            key,
+            state["key"],
             Qs,
             momentum_updates,
             state["balance_counter"],