Pass name of layer to init scale.

PiperOrigin-RevId: 555526825
google · Aug 10, 2023 · 71b71da · 71b71da
1 parent 14ad181
commit 71b71da
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/paxml/learners.py b/paxml/learners.py
@@ -415,7 +415,7 @@ def _adjust_var(old_var, transformed_grad, is_learnable, var_wh, var_key):
             target_magnitude = 1.0
             logging.info('var: %s is a scalar', var_key)
           else:
-            target_magnitude = base_layer.var_init_scale(var_wh)
+            target_magnitude = base_layer.var_init_scale(var_wh, var_key)
           logging.info(
               'var: %s, target_magnitude: %f', var_key, target_magnitude
           )