change call

mosaicml · Feb 6, 2024 · a44d5e8 · a44d5e8
1 parent 105f766
commit a44d5e8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -366,8 +366,8 @@ def _save_checkpoint(self, state: State, logger: Logger):
                                 os.path.join(local_save_path, license_filename),
                             )
 
-                        mlflow_logger.register_model(
+                        mlflow_logger.create_model_version(
                             model_uri=local_save_path,
                             name=self.mlflow_registered_model_name,
-                            await_registration_for=3600,
+                            await_creation_for=3600,
                         )