Fixes named_chain where some kwargs values are configs. (apple#337)

dorucioclea · Feb 21, 2024 · b762f1a · b762f1a
1 parent cf2572a
commit b762f1a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/axlearn/common/optimizers.py b/axlearn/common/optimizers.py
@@ -81,6 +81,8 @@ def partition(param_spec):
 
 
 def named_chain(**kwargs):
+    kwargs = {k: _to_partitioned_transformation(v) for k, v in kwargs.items()}
+
     def init_fn(params):
         return {k: v.init(params) for k, v in kwargs.items()}