From c8bd06f26610aec7a970b23045136b492000cc90 Mon Sep 17 00:00:00 2001
From: Nancy Hung <nancy.hung1068@gmail.com>
Date: Mon, 23 Sep 2024 18:03:23 -0700
Subject: [PATCH] isolated changes

---
 llmfoundry/callbacks/hf_checkpointer.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
index 8956858a93..61d348967c 100644
--- a/llmfoundry/callbacks/hf_checkpointer.py
+++ b/llmfoundry/callbacks/hf_checkpointer.py
@@ -107,7 +107,7 @@ def _maybe_get_license_filename(
         return None
 
 
-def _register_model_multiprocess(
+def _register_model_with_run_id_multiprocess(
     mlflow_logger: MLFlowLogger,
     composer_logging_level: int,
     model_uri: str,
@@ -128,7 +128,7 @@ def _register_model_multiprocess(
         logging.getLogger('composer').setLevel(composer_logging_level)
 
     # Register model.
-    mlflow_logger.register_model(
+    mlflow_logger.register_model_with_run_id(
         model_uri=model_uri,
         name=name,
         await_creation_for=await_creation_for,
@@ -793,7 +793,7 @@ def tensor_hook(
                         process.start()
                     # Faster method to register model in parallel.
                     process = SpawnProcess(
-                        target=_register_model_multiprocess,
+                        target=_register_model_with_run_id_multiprocess,
                         kwargs={
                             'mlflow_logger':
                                 mlflow_logger,