Add task=feature-extraction

huggingface · Nov 22, 2024 · e2828ff · e2828ff
1 parent eeb3159
commit e2828ff
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -2484,6 +2484,13 @@ class PatchTSTOnnxConfig(OnnxConfig):
     def inputs(self) -> Dict[str, Dict[int, str]]:
         return {"past_values": {0: "batch_size", 1: "sequence_length"}}
 
+    @property
+    def outputs(self) -> Dict[str, Dict[int, str]]:
+        if self.task == "feature-extraction":
+            return {"last_hidden_state": {0: "batch_size"}}
+        else:
+            return super().outputs
+
 
 class PatchTSMixerOnnxConfig(PatchTSTOnnxConfig):
     pass