Add logging to modeling tf utils, will be reverted just for debugging

a8nova · May 5, 2024 · 774b6b7 · 774b6b7
1 parent 278938a
commit 774b6b7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/modeling_tf_utils.py b/src/transformers/modeling_tf_utils.py
@@ -2475,7 +2475,7 @@ def save_pretrained(
         output_model_file = os.path.join(save_directory, weights_name)
 
         shards, index = tf_shard_checkpoint(self.weights, max_shard_size, weights_name=weights_name)
-
+        logger.info(f"shards {shards}\nindex {index}")
         # Clean the folder from a previous save
         for filename in os.listdir(save_directory):
             full_filename = os.path.join(save_directory, filename)