From 46ea540ed2fd0bbfe63f08aa6648a77eedf65f84 Mon Sep 17 00:00:00 2001 From: Irene Dea Date: Tue, 3 Dec 2024 23:35:35 +0000 Subject: [PATCH] fix tests --- llmfoundry/callbacks/hf_checkpointer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py index 220e7c53a7..fd6c7aabab 100644 --- a/llmfoundry/callbacks/hf_checkpointer.py +++ b/llmfoundry/callbacks/hf_checkpointer.py @@ -411,7 +411,8 @@ def run_event(self, event: Event, state: State, logger: Logger) -> None: ), upload_to_save_folder=self.save_folder is not None and (not self.final_register_only or not is_last_batch), - register=is_last_batch, # Register only on the last batch + register=self.mlflow_registered_model_name is not None and + is_last_batch, # Register only on the last batch ) elif event == Event.INIT: if not isinstance(state.model, HuggingFaceModel):