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.