We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
MT5 Support 
Big Thanks for the awesome Optimum library, 
MT5 is a Massive Transformers And Supports Many Languages, but it has a slow inference time
for speedup MT5 inference time, We need some post-training methods, and Optimum will help more!

Please Add MT5 Support for Optimum! Thanks

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `optimum/onnxruntime/utils.py`
Content:
```
1 #  Copyright 2021 The HuggingFace Team. All rights reserved.
2 #
3 #  Licensed under the Apache License, Version 2.0 (the "License");
4 #  you may not use this file except in compliance with the License.
5 #  You may obtain a copy of the License at
6 #
7 #      http://www.apache.org/licenses/LICENSE-2.0
8 #
9 #  Unless required by applicable law or agreed to in writing, software
10 #  distributed under the License is distributed on an "AS IS" BASIS,
11 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 #  See the License for the specific language governing permissions and
13 #  limitations under the License.
14 from enum import Enum
15 
16 import torch
17 from transformers.onnx import OnnxConfig, OnnxSeq2SeqConfigWithPast
18 from transformers.utils import logging
19 
20 import onnx
21 import onnxruntime as ort
22 
23 from ..onnx import OnnxConfigWithLoss, OnnxSeq2SeqConfigWithPastAndLoss
24 
25 
26 logger = logging.get_logger(__name__)
27 
28 ONNX_WEIGHTS_NAME = "model.onnx"
29 OPTIMIZED_ONNX_WEIGHTS_NAME = "optimized_model.onnx"
30 QUANTIZED_ONNX_WEIGHTS_NAME = "q8_model.onnx"
31 
32 ONNX_ENCODER_NAME = "encoder_model.onnx"
33 ONNX_DECODER_NAME = "decoder_model.onnx"
34 ONNX_DECODER_WITH_PAST_NAME = "decoder_with_past_model.onnx"
35 
36 
37 def _is_gpu_available():
38     """
39     checks if a gpu is available.
40     """
41     available_providers = ort.get_available_providers()
42     if "CUDAExecutionProvider" in available_providers and torch.cuda.is_available():
43         return True
44     else:
45         return False
46 
47 
48 class ORTConfigManager:
49     """
50     A class that contains all the information needed by ONNX Runtime optimization for a given model type.
51 
52     Attributes:
53         _conf (`Dict[str, tuple]`):
54             A dictionary mapping each supported model type to a tuple containing the number of attention heads
55             and the hidden size model config attribute names as well as the corresponding ONNX Runtime model type.
56     """
57 
58     # Contribution note: Please add new models in alphabetical order
59     _conf = {
60         "albert": ("num_attention_heads", "hidden_size", "bert"),
61         "bart": ("encoder_attention_heads", "d_model", "bart"),
62         "bert": ("num_attention_heads", "hidden_size", "bert"),
63         "big_bird": ("num_attention_heads", "hidden_size", "bert"),
64         "camembert": ("num_attention_heads", "hidden_size", "bert"),
65         "codegen": ("n_head", "n_embd", "gpt2"),
66         "deberta": ("num_attention_heads", "hidden_size", "bert"),
67         "deberta-v2": ("num_attention_heads", "hidden_size", "bert"),
68         "distilbert": ("n_heads", "dim", "bert"),
69         "electra": ("num_attention_heads", "hidden_size", "bert"),
70         "gpt2": ("n_head", "n_embd", "gpt2"),
71         "gpt_neo": ("num_heads", "hidden_size", "gpt2"),
72         "marian": ("encoder_attention_heads", "d_model", "bart"),
73         "roberta": ("num_attention_heads", "hidden_size", "bert"),
74         "xlm-roberta": ("num_attention_heads", "hidden_size", "bert"),
75     }
76 
77     @classmethod
78     def get_num_heads_name(cls, model_type: str) -> str:
79         num_heads = "num_attention_heads"
80         try:
81             num_heads = cls._conf[model_type][0]
82         except KeyError:
83             logger.warning(
84                 f"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to "
85                 f"access the number of heads defined in the config is set to `{num_heads}`."
86             )
87         return num_heads
88 
89     @classmethod
90     def get_hidden_size_name(cls, model_type: str) -> str:
91         hidden_size = "hidden_size"
92         try:
93             hidden_size = cls._conf[model_type][1]
94         except KeyError:
95             logger.warning(
96                 f"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to "
97                 f"access the hidden size defined in the config is set to `{hidden_size}`."
98             )
99         return hidden_size
100 
101     @classmethod
102     def get_model_ort_type(cls, model_type: str) -> str:
103         try:
104             model_type = cls._conf[model_type][2]
105         except KeyError:
106             logger.warning(f"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported.")
107         return model_type
108 
109     @classmethod
110     def check_supported_model_or_raise(cls, model_type: str) -> bool:
111         if model_type not in cls._conf:
112             raise KeyError(
113                 f"{model_type} model type is not supported yet. Only {list(cls._conf.keys())} are supported. "
114                 f"If you want to support {model_type} please propose a PR or open up an issue."
115             )
116 
117 
118 def generate_identified_filename(filename, identifier):
119     return filename.parent.joinpath(filename.stem + identifier).with_suffix(filename.suffix)
120 
121 
122 def fix_atenops_to_gather(model_path):
123     # Fix broken ATenOp nodes back to Gather nodes.
124     model = onnx.load(model_path)
125     onnx.checker.check_model(model)
126 
127     nodes = model.graph.node
128 
129     for node in nodes:
130         if node.op_type in ["ATenOp", "ATen"]:
131             logger.info(f"----Start fixing node: {node.name}----")
132             op_num = node.name.split("_")[-1]
133             new_node = onnx.helper.make_node(
134                 "Gather",
135                 name="Gather_" + op_num,
136                 inputs=[node.input[0], node.input[1]],
137                 outputs=node.output,
138             )
139 
140             model.graph.node.remove(node)
141             model.graph.node.insert(int(op_num), new_node)
142 
143     onnx.checker.check_model(model)
144     onnx.save(model, model_path)
145 
146 
147 def wrap_onnx_config_for_loss(onnx_config: OnnxConfig) -> OnnxConfig:
148     if isinstance(onnx_config, OnnxSeq2SeqConfigWithPast):
149         return OnnxSeq2SeqConfigWithPastAndLoss(onnx_config)
150     else:
151         return OnnxConfigWithLoss(onnx_config)
152 
153 
154 def get_device_for_provider(provider: str) -> torch.device:
155     """
156     Gets the PyTorch device (CPU/CUDA) associated with an ONNX Runtime provider.
157     """
158     return torch.device("cuda") if provider == "CUDAExecutionProvider" else torch.device("cpu")
159 
160 
161 def get_provider_for_device(device: torch.device) -> str:
162     """
163     Gets the ONNX Runtime provider associated with the PyTorch device (CPU/CUDA).
164     """
165     return "CUDAExecutionProvider" if device.type.lower() == "cuda" else "CPUExecutionProvider"
166 
167 
168 class ORTQuantizableOperator(Enum):
169     # Common ops
170     Gather = "Gather"
171     Transpose = "Transpose"
172     EmbedLayerNormalizationQuant = "EmbedLayerNormalization"
173 
174     # QLinearOps
175     Conv = "Conv"
176     MatMul = "MatMul"
177     Add = "Add"
178     Mul = "Mul"
179     Relu = "Relu"
180     Clip = "Clip"
181     LeakyRelu = "LeakyRelu"
182     Sigmoid = "Sigmoid"
183     MaxPool = "MaxPool"
184     GlobalAveragePool = "GlobalAveragePool"
185     Split = "Split"
186     Pad = "Pad"
187     Reshape = "Reshape"
188     Squeeze = "Squeeze"
189     Unsqueeze = "Unsqueeze"
190     Resize = "Resize"
191     AveragePool = "AveragePool"
192     Concat = "Concat"
193
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py
--- a/optimum/onnxruntime/utils.py
+++ b/optimum/onnxruntime/utils.py
@@ -69,6 +69,7 @@
         "electra": ("num_attention_heads", "hidden_size", "bert"),
         "gpt2": ("n_head", "n_embd", "gpt2"),
         "gpt_neo": ("num_heads", "hidden_size", "gpt2"),
+        "mt5": ("num_heads", "d_model", "bart"),
         "marian": ("encoder_attention_heads", "d_model", "bart"),
         "roberta": ("num_attention_heads", "hidden_size", "bert"),
         "xlm-roberta": ("num_attention_heads", "hidden_size", "bert"),

{"golden_diff": "diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py\n--- a/optimum/onnxruntime/utils.py\n+++ b/optimum/onnxruntime/utils.py\n@@ -69,6 +69,7 @@\n         \"electra\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n         \"gpt2\": (\"n_head\", \"n_embd\", \"gpt2\"),\n         \"gpt_neo\": (\"num_heads\", \"hidden_size\", \"gpt2\"),\n+        \"mt5\": (\"num_heads\", \"d_model\", \"bart\"),\n         \"marian\": (\"encoder_attention_heads\", \"d_model\", \"bart\"),\n         \"roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n         \"xlm-roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n", "issue": "MT5 Support \nBig Thanks for the awesome Optimum library, \r\nMT5 is a Massive Transformers And Supports Many Languages, but it has a slow inference time\r\nfor speedup MT5 inference time, We need some post-training methods, and Optimum will help more!\r\n\r\nPlease Add MT5 Support for Optimum! Thanks\n", "before_files": [{"content": "#  Copyright 2021 The HuggingFace Team. All rights reserved.\n#\n#  Licensed under the Apache License, Version 2.0 (the \"License\");\n#  you may not use this file except in compliance with the License.\n#  You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n#  Unless required by applicable law or agreed to in writing, software\n#  distributed under the License is distributed on an \"AS IS\" BASIS,\n#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n#  See the License for the specific language governing permissions and\n#  limitations under the License.\nfrom enum import Enum\n\nimport torch\nfrom transformers.onnx import OnnxConfig, OnnxSeq2SeqConfigWithPast\nfrom transformers.utils import logging\n\nimport onnx\nimport onnxruntime as ort\n\nfrom ..onnx import OnnxConfigWithLoss, OnnxSeq2SeqConfigWithPastAndLoss\n\n\nlogger = logging.get_logger(__name__)\n\nONNX_WEIGHTS_NAME = \"model.onnx\"\nOPTIMIZED_ONNX_WEIGHTS_NAME = \"optimized_model.onnx\"\nQUANTIZED_ONNX_WEIGHTS_NAME = \"q8_model.onnx\"\n\nONNX_ENCODER_NAME = \"encoder_model.onnx\"\nONNX_DECODER_NAME = \"decoder_model.onnx\"\nONNX_DECODER_WITH_PAST_NAME = \"decoder_with_past_model.onnx\"\n\n\ndef _is_gpu_available():\n    \"\"\"\n    checks if a gpu is available.\n    \"\"\"\n    available_providers = ort.get_available_providers()\n    if \"CUDAExecutionProvider\" in available_providers and torch.cuda.is_available():\n        return True\n    else:\n        return False\n\n\nclass ORTConfigManager:\n    \"\"\"\n    A class that contains all the information needed by ONNX Runtime optimization for a given model type.\n\n    Attributes:\n        _conf (`Dict[str, tuple]`):\n            A dictionary mapping each supported model type to a tuple containing the number of attention heads\n            and the hidden size model config attribute names as well as the corresponding ONNX Runtime model type.\n    \"\"\"\n\n    # Contribution note: Please add new models in alphabetical order\n    _conf = {\n        \"albert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"bart\": (\"encoder_attention_heads\", \"d_model\", \"bart\"),\n        \"bert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"big_bird\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"camembert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"codegen\": (\"n_head\", \"n_embd\", \"gpt2\"),\n        \"deberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"deberta-v2\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"distilbert\": (\"n_heads\", \"dim\", \"bert\"),\n        \"electra\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"gpt2\": (\"n_head\", \"n_embd\", \"gpt2\"),\n        \"gpt_neo\": (\"num_heads\", \"hidden_size\", \"gpt2\"),\n        \"marian\": (\"encoder_attention_heads\", \"d_model\", \"bart\"),\n        \"roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"xlm-roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n    }\n\n    @classmethod\n    def get_num_heads_name(cls, model_type: str) -> str:\n        num_heads = \"num_attention_heads\"\n        try:\n            num_heads = cls._conf[model_type][0]\n        except KeyError:\n            logger.warning(\n                f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to \"\n                f\"access the number of heads defined in the config is set to `{num_heads}`.\"\n            )\n        return num_heads\n\n    @classmethod\n    def get_hidden_size_name(cls, model_type: str) -> str:\n        hidden_size = \"hidden_size\"\n        try:\n            hidden_size = cls._conf[model_type][1]\n        except KeyError:\n            logger.warning(\n                f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to \"\n                f\"access the hidden size defined in the config is set to `{hidden_size}`.\"\n            )\n        return hidden_size\n\n    @classmethod\n    def get_model_ort_type(cls, model_type: str) -> str:\n        try:\n            model_type = cls._conf[model_type][2]\n        except KeyError:\n            logger.warning(f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported.\")\n        return model_type\n\n    @classmethod\n    def check_supported_model_or_raise(cls, model_type: str) -> bool:\n        if model_type not in cls._conf:\n            raise KeyError(\n                f\"{model_type} model type is not supported yet. Only {list(cls._conf.keys())} are supported. \"\n                f\"If you want to support {model_type} please propose a PR or open up an issue.\"\n            )\n\n\ndef generate_identified_filename(filename, identifier):\n    return filename.parent.joinpath(filename.stem + identifier).with_suffix(filename.suffix)\n\n\ndef fix_atenops_to_gather(model_path):\n    # Fix broken ATenOp nodes back to Gather nodes.\n    model = onnx.load(model_path)\n    onnx.checker.check_model(model)\n\n    nodes = model.graph.node\n\n    for node in nodes:\n        if node.op_type in [\"ATenOp\", \"ATen\"]:\n            logger.info(f\"----Start fixing node: {node.name}----\")\n            op_num = node.name.split(\"_\")[-1]\n            new_node = onnx.helper.make_node(\n                \"Gather\",\n                name=\"Gather_\" + op_num,\n                inputs=[node.input[0], node.input[1]],\n                outputs=node.output,\n            )\n\n            model.graph.node.remove(node)\n            model.graph.node.insert(int(op_num), new_node)\n\n    onnx.checker.check_model(model)\n    onnx.save(model, model_path)\n\n\ndef wrap_onnx_config_for_loss(onnx_config: OnnxConfig) -> OnnxConfig:\n    if isinstance(onnx_config, OnnxSeq2SeqConfigWithPast):\n        return OnnxSeq2SeqConfigWithPastAndLoss(onnx_config)\n    else:\n        return OnnxConfigWithLoss(onnx_config)\n\n\ndef get_device_for_provider(provider: str) -> torch.device:\n    \"\"\"\n    Gets the PyTorch device (CPU/CUDA) associated with an ONNX Runtime provider.\n    \"\"\"\n    return torch.device(\"cuda\") if provider == \"CUDAExecutionProvider\" else torch.device(\"cpu\")\n\n\ndef get_provider_for_device(device: torch.device) -> str:\n    \"\"\"\n    Gets the ONNX Runtime provider associated with the PyTorch device (CPU/CUDA).\n    \"\"\"\n    return \"CUDAExecutionProvider\" if device.type.lower() == \"cuda\" else \"CPUExecutionProvider\"\n\n\nclass ORTQuantizableOperator(Enum):\n    # Common ops\n    Gather = \"Gather\"\n    Transpose = \"Transpose\"\n    EmbedLayerNormalizationQuant = \"EmbedLayerNormalization\"\n\n    # QLinearOps\n    Conv = \"Conv\"\n    MatMul = \"MatMul\"\n    Add = \"Add\"\n    Mul = \"Mul\"\n    Relu = \"Relu\"\n    Clip = \"Clip\"\n    LeakyRelu = \"LeakyRelu\"\n    Sigmoid = \"Sigmoid\"\n    MaxPool = \"MaxPool\"\n    GlobalAveragePool = \"GlobalAveragePool\"\n    Split = \"Split\"\n    Pad = \"Pad\"\n    Reshape = \"Reshape\"\n    Squeeze = \"Squeeze\"\n    Unsqueeze = \"Unsqueeze\"\n    Resize = \"Resize\"\n    AveragePool = \"AveragePool\"\n    Concat = \"Concat\"\n", "path": "optimum/onnxruntime/utils.py"}], "after_files": [{"content": "#  Copyright 2021 The HuggingFace Team. All rights reserved.\n#\n#  Licensed under the Apache License, Version 2.0 (the \"License\");\n#  you may not use this file except in compliance with the License.\n#  You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n#  Unless required by applicable law or agreed to in writing, software\n#  distributed under the License is distributed on an \"AS IS\" BASIS,\n#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n#  See the License for the specific language governing permissions and\n#  limitations under the License.\nfrom enum import Enum\n\nimport torch\nfrom transformers.onnx import OnnxConfig, OnnxSeq2SeqConfigWithPast\nfrom transformers.utils import logging\n\nimport onnx\nimport onnxruntime as ort\n\nfrom ..onnx import OnnxConfigWithLoss, OnnxSeq2SeqConfigWithPastAndLoss\n\n\nlogger = logging.get_logger(__name__)\n\nONNX_WEIGHTS_NAME = \"model.onnx\"\nOPTIMIZED_ONNX_WEIGHTS_NAME = \"optimized_model.onnx\"\nQUANTIZED_ONNX_WEIGHTS_NAME = \"q8_model.onnx\"\n\nONNX_ENCODER_NAME = \"encoder_model.onnx\"\nONNX_DECODER_NAME = \"decoder_model.onnx\"\nONNX_DECODER_WITH_PAST_NAME = \"decoder_with_past_model.onnx\"\n\n\ndef _is_gpu_available():\n    \"\"\"\n    checks if a gpu is available.\n    \"\"\"\n    available_providers = ort.get_available_providers()\n    if \"CUDAExecutionProvider\" in available_providers and torch.cuda.is_available():\n        return True\n    else:\n        return False\n\n\nclass ORTConfigManager:\n    \"\"\"\n    A class that contains all the information needed by ONNX Runtime optimization for a given model type.\n\n    Attributes:\n        _conf (`Dict[str, tuple]`):\n            A dictionary mapping each supported model type to a tuple containing the number of attention heads\n            and the hidden size model config attribute names as well as the corresponding ONNX Runtime model type.\n    \"\"\"\n\n    # Contribution note: Please add new models in alphabetical order\n    _conf = {\n        \"albert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"bart\": (\"encoder_attention_heads\", \"d_model\", \"bart\"),\n        \"bert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"big_bird\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"camembert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"codegen\": (\"n_head\", \"n_embd\", \"gpt2\"),\n        \"deberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"deberta-v2\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"distilbert\": (\"n_heads\", \"dim\", \"bert\"),\n        \"electra\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"gpt2\": (\"n_head\", \"n_embd\", \"gpt2\"),\n        \"gpt_neo\": (\"num_heads\", \"hidden_size\", \"gpt2\"),\n        \"mt5\": (\"num_heads\", \"d_model\", \"bart\"),\n        \"marian\": (\"encoder_attention_heads\", \"d_model\", \"bart\"),\n        \"roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"xlm-roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n    }\n\n    @classmethod\n    def get_num_heads_name(cls, model_type: str) -> str:\n        num_heads = \"num_attention_heads\"\n        try:\n            num_heads = cls._conf[model_type][0]\n        except KeyError:\n            logger.warning(\n                f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to \"\n                f\"access the number of heads defined in the config is set to `{num_heads}`.\"\n            )\n        return num_heads\n\n    @classmethod\n    def get_hidden_size_name(cls, model_type: str) -> str:\n        hidden_size = \"hidden_size\"\n        try:\n            hidden_size = cls._conf[model_type][1]\n        except KeyError:\n            logger.warning(\n                f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to \"\n                f\"access the hidden size defined in the config is set to `{hidden_size}`.\"\n            )\n        return hidden_size\n\n    @classmethod\n    def get_model_ort_type(cls, model_type: str) -> str:\n        try:\n            model_type = cls._conf[model_type][2]\n        except KeyError:\n            logger.warning(f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported.\")\n        return model_type\n\n    @classmethod\n    def check_supported_model_or_raise(cls, model_type: str) -> bool:\n        if model_type not in cls._conf:\n            raise KeyError(\n                f\"{model_type} model type is not supported yet. Only {list(cls._conf.keys())} are supported. \"\n                f\"If you want to support {model_type} please propose a PR or open up an issue.\"\n            )\n\n\ndef generate_identified_filename(filename, identifier):\n    return filename.parent.joinpath(filename.stem + identifier).with_suffix(filename.suffix)\n\n\ndef fix_atenops_to_gather(model_path):\n    # Fix broken ATenOp nodes back to Gather nodes.\n    model = onnx.load(model_path)\n    onnx.checker.check_model(model)\n\n    nodes = model.graph.node\n\n    for node in nodes:\n        if node.op_type in [\"ATenOp\", \"ATen\"]:\n            logger.info(f\"----Start fixing node: {node.name}----\")\n            op_num = node.name.split(\"_\")[-1]\n            new_node = onnx.helper.make_node(\n                \"Gather\",\n                name=\"Gather_\" + op_num,\n                inputs=[node.input[0], node.input[1]],\n                outputs=node.output,\n            )\n\n            model.graph.node.remove(node)\n            model.graph.node.insert(int(op_num), new_node)\n\n    onnx.checker.check_model(model)\n    onnx.save(model, model_path)\n\n\ndef wrap_onnx_config_for_loss(onnx_config: OnnxConfig) -> OnnxConfig:\n    if isinstance(onnx_config, OnnxSeq2SeqConfigWithPast):\n        return OnnxSeq2SeqConfigWithPastAndLoss(onnx_config)\n    else:\n        return OnnxConfigWithLoss(onnx_config)\n\n\ndef get_device_for_provider(provider: str) -> torch.device:\n    \"\"\"\n    Gets the PyTorch device (CPU/CUDA) associated with an ONNX Runtime provider.\n    \"\"\"\n    return torch.device(\"cuda\") if provider == \"CUDAExecutionProvider\" else torch.device(\"cpu\")\n\n\ndef get_provider_for_device(device: torch.device) -> str:\n    \"\"\"\n    Gets the ONNX Runtime provider associated with the PyTorch device (CPU/CUDA).\n    \"\"\"\n    return \"CUDAExecutionProvider\" if device.type.lower() == \"cuda\" else \"CPUExecutionProvider\"\n\n\nclass ORTQuantizableOperator(Enum):\n    # Common ops\n    Gather = \"Gather\"\n    Transpose = \"Transpose\"\n    EmbedLayerNormalizationQuant = \"EmbedLayerNormalization\"\n\n    # QLinearOps\n    Conv = \"Conv\"\n    MatMul = \"MatMul\"\n    Add = \"Add\"\n    Mul = \"Mul\"\n    Relu = \"Relu\"\n    Clip = \"Clip\"\n    LeakyRelu = \"LeakyRelu\"\n    Sigmoid = \"Sigmoid\"\n    MaxPool = \"MaxPool\"\n    GlobalAveragePool = \"GlobalAveragePool\"\n    Split = \"Split\"\n    Pad = \"Pad\"\n    Reshape = \"Reshape\"\n    Squeeze = \"Squeeze\"\n    Unsqueeze = \"Unsqueeze\"\n    Resize = \"Resize\"\n    AveragePool = \"AveragePool\"\n    Concat = \"Concat\"\n", "path": "optimum/onnxruntime/utils.py"}]}