Fix OV trainer model saving step (#472)

huggingface · Nov 6, 2023 · 9562235 · 9562235
1 parent 59f63a5
commit 9562235
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)