update the logger message with accordant weights_file_name (#28181)

Co-authored-by: yudong.lin <[email protected]>
huggingface · Dec 22, 2023 · 68fa1e8 · 68fa1e8
1 parent 74d9d0c
commit 68fa1e8
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
@@ -2408,7 +2408,8 @@ def save_pretrained(
                 save_function(shard, os.path.join(save_directory, shard_file))
 
         if index is None:
-            path_to_weights = os.path.join(save_directory, _add_variant(WEIGHTS_NAME, variant))
+            weights_file_name = SAFE_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME
+            path_to_weights = os.path.join(save_directory, _add_variant(weights_file_name, variant))
             logger.info(f"Model weights saved in {path_to_weights}")
         else:
             save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME