fix output names for maskformer export

huggingface · Dec 18, 2024 · 3fa346c · 3fa346c
1 parent 01929b2
commit 3fa346c
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -1001,6 +1001,12 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
         else:
             return super().outputs
 
+    @property
+    def torch_to_onnx_output_map(self) -> Dict[str, str]:
+        return {
+            "transformer_decoder_last_hidden_state": "last_hidden_state",
+        }
+
 
 class DonutSwinOnnxConfig(ViTOnnxConfig):
     DEFAULT_ONNX_OPSET = 11