diff --git a/src/transformers/onnx/features.py b/src/transformers/onnx/features.py index e7c24a8ad97a81..3eea94c8c1a64e 100644 --- a/src/transformers/onnx/features.py +++ b/src/transformers/onnx/features.py @@ -229,6 +229,12 @@ class FeaturesManager: "question-answering", onnx_config_cls="models.data2vec.Data2VecTextOnnxConfig", ), + "data2vec-vision": supported_features_mapping( + "default", + "image-classification", + "image-segmentation", + onnx_config_cls="models.data2vec.Data2VecVisionOnnxConfig", + ), "deberta": supported_features_mapping( "default", "masked-lm", diff --git a/tests/onnx/test_onnx_v2.py b/tests/onnx/test_onnx_v2.py index c15910734f3adb..50fc3f59c59e4b 100644 --- a/tests/onnx/test_onnx_v2.py +++ b/tests/onnx/test_onnx_v2.py @@ -207,6 +207,7 @@ def test_values_override(self): ("deit", "facebook/deit-small-patch16-224"), ("beit", "microsoft/beit-base-patch16-224"), ("data2vec-text", "facebook/data2vec-text-base"), + ("data2vec-vision", "facebook/data2vec-vision-base"), ("perceiver", "deepmind/language-perceiver", ("masked-lm", "sequence-classification")), ("perceiver", "deepmind/vision-perceiver-conv", ("image-classification",)), ("yolos", "hustvl/yolos-tiny"),