From bfa8fdb910c5b93a806aed3fb9964dca04207271 Mon Sep 17 00:00:00 2001
From: David Corvoysier <david@huggingface.co>
Date: Mon, 15 Apr 2024 07:27:08 +0000
Subject: [PATCH] feat(decoder): do not split checkpoint files

---
 optimum/neuron/modeling_decoder.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/optimum/neuron/modeling_decoder.py b/optimum/neuron/modeling_decoder.py
index ffdf9d87b..8268c0514 100644
--- a/optimum/neuron/modeling_decoder.py
+++ b/optimum/neuron/modeling_decoder.py
@@ -42,7 +42,6 @@
 
 if is_transformers_neuronx_available():
     from transformers_neuronx.config import ContinuousBatchingConfig, NeuronConfig
-    from transformers_neuronx.module import save_split
 
 
 if TYPE_CHECKING:
@@ -249,9 +248,7 @@ def _create_checkpoint(
 
         # Save the model checkpoint in a temporary directory
         checkpoint_dir = TemporaryDirectory()
-        model.save_pretrained(
-            checkpoint_dir.name, save_function=save_split, safe_serialization=False, max_shard_size="10000GB"
-        )
+        model.save_pretrained(checkpoint_dir.name)
         return checkpoint_dir
 
     @classmethod