add kwargs

stanford-crfm · May 1, 2024 · 828b5d9 · 828b5d9
1 parent b942d60
commit 828b5d9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/optim/schedulefree_adam.py b/src/levanter/optim/schedulefree_adam.py
@@ -98,7 +98,7 @@ def init_fn(params):
             nu=nu,
         )
 
-    def update_fn(updates, state, params=None):
+    def update_fn(updates, state, params=None, **kwargs):
         z = state.z
         t = state.count
         weight_sum = state.weight_sum