From 44ecbf27f98321bd256c4936332fee776c86db90 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Fri, 3 Nov 2023 19:27:48 +0100 Subject: [PATCH] fix test --- optimum/intel/openvino/trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/optimum/intel/openvino/trainer.py b/optimum/intel/openvino/trainer.py index 17b0aa7058..dfc659882c 100644 --- 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)