Fix Yolos ONNX export test (huggingface#18606)

Co-authored-by: lewtun <[email protected]>
Co-authored-by: ydshieh <[email protected]>

Author: 2 people

tests/onnx/test_onnx_v2.py

@@ -284,6 +284,12 @@ def _onnx_export(self, test_name, name, model_name, feature, onnx_config_class_c
         model_class = FeaturesManager.get_model_class_for_feature(feature)
         config = AutoConfig.from_pretrained(model_name)
         model = model_class.from_config(config)
+
+        # Dynamic axes aren't supported for YOLO-like models. This means they cannot be exported to ONNX on CUDA devices.
+        # See: https://github.com/ultralytics/yolov5/pull/8378
+        if model.__class__.__name__.startswith("Yolos") and device != "cpu":
+            return
+
         onnx_config = onnx_config_class_constructor(model.config)
 
         if is_torch_available():