Fixed test issue

huggingface · Nov 10, 2023 · 0cb7f19 · 0cb7f19
1 parent 7815333
commit 0cb7f19
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/optimum/intel/openvino/quantization.py b/optimum/intel/openvino/quantization.py
@@ -41,8 +41,8 @@
 from ..utils.constant import _TASK_ALIASES
 from .configuration import OVConfig
 from .modeling_base import OVBaseModel
+from .modeling_base_seq2seq import OVBaseModelForSeq2SeqLM
 from .modeling_decoder import OVBaseDecoderModel
-from .modeling_seq2seq import OVModelForSeq2SeqLM
 from .utils import (
     MAX_ONNX_OPSET,
     MIN_ONNX_QDQ_OPSET,
@@ -177,23 +177,23 @@ def quantize(
                     "In case you only want to apply quantization on the weights, please set `weights_only=True`."
                 )
 
-        if isinstance(self.model, OVBaseDecoderModel) and self.model.use_cache:
-            self._quantize_ovcausallm(
+        if isinstance(self.model, OVBaseModelForSeq2SeqLM):
+            self._quantize_ovmodelforseq2seqlm(
                 calibration_dataset,
                 save_directory,
                 batch_size,
                 data_collator,
                 remove_unused_columns,
-                weights_only,
                 **kwargs,
             )
-        if isinstance(self.model, OVModelForSeq2SeqLM):
-            self._quantize_ovmodelforseq2seqlm(
+        if isinstance(self.model, OVBaseDecoderModel) and self.model.use_cache:
+            self._quantize_ovcausallm(
                 calibration_dataset,
                 save_directory,
                 batch_size,
                 data_collator,
                 remove_unused_columns,
+                weights_only,
                 **kwargs,
             )
         elif isinstance(self.model, OVBaseModel):