Add SegFormer ONNX support (#18006)

* Add ONNX support * Make height and width dynamic axes Co-authored-by: Niels Rogge <nielsrogge@Nielss-MacBook-Pro.local>
2022-08-31 20:58:44 +02:00
parent 89514f0541
commit 7e7f743481
5 changed files with 38 additions and 2 deletions
--- a/tests/onnx/test_onnx_v2.py
+++ b/tests/onnx/test_onnx_v2.py
@@ -216,6 +216,7 @@ PYTORCH_EXPORT_MODELS = {
    ("perceiver", "deepmind/vision-perceiver-conv", ("image-classification",)),
    ("longformer", "allenai/longformer-base-4096"),
    ("yolos", "hustvl/yolos-tiny"),
+    ("segformer", "nvidia/segformer-b0-finetuned-ade-512-512"),
 }

 PYTORCH_EXPORT_WITH_PAST_MODELS = {