From a177b1e307655b63325707b75c292a6e407d246f Mon Sep 17 00:00:00 2001
From: Irene Dea <deaairene@gmail.com>
Date: Wed, 14 Aug 2024 10:51:33 -0700
Subject: [PATCH] Add some more debug logs

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

diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
index 79dc73de98..78cd99be90 100644
--- a/llmfoundry/callbacks/hf_checkpointer.py
+++ b/llmfoundry/callbacks/hf_checkpointer.py
@@ -125,6 +125,8 @@ def _register_model_with_run_id_multiprocess(
         )
         logging.getLogger('composer').setLevel(composer_logging_level)
 
+    log.debug('Starting registering model.')
+
     # Register model.
     mlflow_logger.register_model_with_run_id(
         model_uri=model_uri,
@@ -132,6 +134,8 @@ def _register_model_with_run_id_multiprocess(
         await_creation_for=await_creation_for,
     )
 
+    log.debug('Finished registering model.')
+
 
 class HuggingFaceCheckpointer(Callback):
     """Save a huggingface formatted checkpoint during training.