From 8259e834c3ac340cf661014f58b9f76a50b835ce Mon Sep 17 00:00:00 2001
From: Ella Charlaix <ella@huggingface.co>
Date: Tue, 17 Oct 2023 18:02:30 +0200
Subject: [PATCH] fix

---
 optimum/intel/neural_compressor/modeling_base.py | 5 +++--
 tests/neural_compressor/test_modeling.py         | 8 +++-----
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/optimum/intel/neural_compressor/modeling_base.py b/optimum/intel/neural_compressor/modeling_base.py
index 8746a8d817..e6ae0f2595 100644
--- a/optimum/intel/neural_compressor/modeling_base.py
+++ b/optimum/intel/neural_compressor/modeling_base.py
@@ -112,6 +112,7 @@ def _from_pretrained(
         file_name: str = WEIGHTS_NAME,
         local_files_only: bool = False,
         subfolder: str = "",
+        trust_remote_code: bool = False,
         **kwargs,
     ):
         model_name_or_path = kwargs.pop("model_name_or_path", None)
@@ -178,8 +179,8 @@ def _from_pretrained(
             model, config=config, model_save_dir=model_save_dir, q_config=q_config, inc_config=inc_config, **kwargs
         )
 
-    def _save_pretrained(self, save_directory: Union[str, Path], file_name: str = WEIGHTS_NAME):
-        output_path = os.path.join(save_directory, file_name)
+    def _save_pretrained(self, save_directory: Union[str, Path]):
+        output_path = os.path.join(save_directory, WEIGHTS_NAME)
 
         if isinstance(self.model, torch.nn.Module):
             state_dict = self.model.state_dict()
diff --git a/tests/neural_compressor/test_modeling.py b/tests/neural_compressor/test_modeling.py
index 9d2f07f0bd..74e75bc666 100644
--- a/tests/neural_compressor/test_modeling.py
+++ b/tests/neural_compressor/test_modeling.py
@@ -38,8 +38,7 @@
     INCStableDiffusionPipeline,
     INCTrainer,
 )
-from optimum.intel.neural_compressor.utils import _HEAD_TO_AUTOMODELS
-
+from optimum.intel.neural_compressor.utils import _HEAD_TO_AUTOMODELS, WEIGHTS_NAME
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 set_seed(1009)
@@ -94,10 +93,9 @@ def test_compare_to_transformers(self, model_id, task):
         config = config_class(inc_model.config)
         model_inputs = config.generate_dummy_inputs(framework="pt")
         outputs = inc_model(**model_inputs)
-        file_name = "model.pt"
         with tempfile.TemporaryDirectory() as tmpdirname:
-            inc_model.save_pretrained(tmpdirname, file_name)
-            loaded_model = model_class.from_pretrained(tmpdirname, file_name=file_name)
+            inc_model.save_pretrained(tmpdirname)
+            loaded_model = model_class.from_pretrained(tmpdirname, file_name=WEIGHTS_NAME)
             outputs_loaded = loaded_model(**model_inputs)
 
         if task == "feature-extraction":