From 2782292fdc5a12d21411f6db48a608bbfd0ae0b1 Mon Sep 17 00:00:00 2001 From: Zach Mueller Date: Mon, 29 Apr 2024 14:47:26 -0400 Subject: [PATCH] Include safetensors as part of `_load_best_model` (#30553) * Include safetensors * Cleanup --- src/transformers/trainer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 1d3c164984ea1c..d967c9314b1025 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2611,7 +2611,9 @@ def _load_best_model(self): load_result = model.load_state_dict(state_dict, False) if not is_sagemaker_mp_enabled() and has_been_loaded: self._issue_warnings_after_load(load_result) - elif os.path.exists(os.path.join(self.state.best_model_checkpoint, WEIGHTS_INDEX_NAME)): + elif os.path.exists(os.path.join(self.state.best_model_checkpoint, SAFE_WEIGHTS_INDEX_NAME)) or os.path.exists( + os.path.join(self.state.best_model_checkpoint, WEIGHTS_INDEX_NAME) + ): load_result = load_sharded_checkpoint( model, self.state.best_model_checkpoint, strict=is_sagemaker_mp_enabled() )