better custom config

Author: IlyasMoutawwakil

tests/openvino/test_export.py

@@ -124,17 +124,21 @@ def test_export_custom_model(self):
         class BertOnnxConfigWithPooler(BertOnnxConfig):
             @property
             def outputs(self):
-                common_outputs = {}
-                common_outputs["last_hidden_state"] = {0: "batch_size", 1: "sequence_length"}
-                common_outputs["pooler_output"] = {0: "batch_size"}
+                if self.task == "feature-extraction-with-pooler":
+                    common_outputs = {}
+                    common_outputs["last_hidden_state"] = {0: "batch_size", 1: "sequence_length"}
+                    common_outputs["pooler_output"] = {0: "batch_size"}
+                else:
+                    common_outputs = super().outputs
+
                 return common_outputs
 
         base_task = "feature-extraction"
         custom_task = f"{base_task}-with-pooler"
         model_id = "sentence-transformers/all-MiniLM-L6-v2"
 
         config = AutoConfig.from_pretrained(model_id)
-        custom_export_configs = {"model": BertOnnxConfigWithPooler(config, task=base_task)}
+        custom_export_configs = {"model": BertOnnxConfigWithPooler(config, task=custom_task)}
 
         with TemporaryDirectory() as tmpdirname:
             main_export(