From 3fa346c21f09c9333663130c784a8f813fa3ac8c Mon Sep 17 00:00:00 2001
From: Ella Charlaix <ella@huggingface.co>
Date: Wed, 18 Dec 2024 17:37:28 +0100
Subject: [PATCH] fix output names for maskformer export

---
 optimum/exporters/onnx/model_configs.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
index 771c9de70f..4c5a727a18 100644
--- 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