fix test

huggingface · Nov 3, 2023 · 44ecbf2 · 44ecbf2
1 parent 59f63a5
commit 44ecbf2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/optimum/intel/openvino/trainer.py b/optimum/intel/openvino/trainer.py
@@ -827,12 +827,12 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None):
             if state_dict is None:
                 state_dict = self.model.state_dict()
             if is_pretrained_model:
-                unwrapped_model.save_pretrained(output_dir, state_dict=state_dict)
+                unwrapped_model.save_pretrained(output_dir, state_dict=state_dict, safe_serialization=False)
             else:
                 logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                 torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
         else:
-            self.model.save_pretrained(output_dir, state_dict=state_dict)
+            self.model.save_pretrained(output_dir, state_dict=state_dict, safe_serialization=False)
 
         if self.tokenizer is not None:
             self.tokenizer.save_pretrained(output_dir)