Revert mutable kwarg in abstract_init_with_metadata in init checkpoin…

…t rule
google · Apr 10, 2024 · ed3300f · ed3300f
1 parent f94e783
commit ed3300f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/paxml/tasks_lib.py b/paxml/tasks_lib.py
@@ -1787,7 +1787,7 @@ def _apply_init_checkpoint_rule(
       )
     # Initialize with a dummy seed
     var_weight_hparams = ckpt_task.model.abstract_init_with_metadata(
-        inputs_shape_dtype, mutable=DEFAULT_INIT_MUTABLE_LIST)
+        inputs_shape_dtype, extra_mutable_list=DEFAULT_INIT_MUTABLE_LIST)
     ckpt_train_state = ckpt_task.create_train_state_padded_shapes(
         var_weight_hparams)
     train_state_pspecs = ckpt_task.create_train_state_partition_specs(