diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py index 0b0edd04a99..321f18ad692 100644 --- a/optimum/exporters/onnx/model_configs.py +++ b/optimum/exporters/onnx/model_configs.py @@ -117,6 +117,7 @@ def inputs(self) -> Dict[str, Dict[int, str]]: "token_type_ids": dynamic_axis, } + class VisualBertOnnxConfig(TextAndVisionOnnxConfig): DEFAULT_ONNX_OPSET = 11 @@ -127,7 +128,7 @@ def inputs(self) -> Dict[str, Dict[int, str]]: "attention_mask": {0: "batch_size", 1: "sequence_length"}, "pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, } - + @property def outputs(self) -> Dict[str, Dict[int, str]]: return { diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py index 87cbed4f331..b8160d65673 100644 --- a/tests/exporters/exporters_utils.py +++ b/tests/exporters/exporters_utils.py @@ -197,7 +197,7 @@ "document-question-answering-with-past", ], }, - "visualbert": "hf-internal-testing/tiny-random-VisualBertModel", + "visualbert": "hf-internal-testing/tiny-random-VisualBertModel", }