diff --git a/tests/onnx/test_onnx_v2.py b/tests/onnx/test_onnx_v2.py index 5634abc77068..2055bd49a7c3 100644 --- a/tests/onnx/test_onnx_v2.py +++ b/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():