From c99e984657b64dd8f19de74405bbf13763ab4f2b Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Wed, 17 Aug 2022 10:04:49 +0200 Subject: [PATCH] Fix Yolos ONNX export test (#18606) Co-authored-by: lewtun Co-authored-by: ydshieh --- tests/onnx/test_onnx_v2.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/onnx/test_onnx_v2.py b/tests/onnx/test_onnx_v2.py index 79eff60cef..b3c0ffb1f3 100644 --- a/tests/onnx/test_onnx_v2.py +++ b/tests/onnx/test_onnx_v2.py @@ -284,6 +284,12 @@ class OnnxExportTestCaseV2(TestCase): 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():