Update kron.py

stanford-crfm · Dec 19, 2024 · b62963e · b62963e
1 parent 3fa70ab
commit b62963e
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/levanter/optim/kron.py b/src/levanter/optim/kron.py
@@ -133,9 +133,10 @@ def _optimizer(learning_rate) -> optax.GradientTransformation:
             components.append(optax.scale_by_learning_rate(learning_rate))
             return optax.chain(*components)
 
-        return optax.inject_hyperparams(_optimizer)(
-            learning_rate=self.lr_scheduler(num_train_steps)
-        )
+        # return optax.inject_hyperparams(_optimizer)(
+        #     learning_rate=self.lr_scheduler(num_train_steps)
+        # )
+        return _optimizer(self.lr_scheduler(num_train_steps))
 
 
 from typing import Any, List, Optional, Union, Callable