Add semantic-segmentation to config

huggingface · Jul 2, 2022 · b3201eb · b3201eb
1 parent 284b6e1
commit b3201eb
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/transformers/onnx/config.py b/src/transformers/onnx/config.py
@@ -99,6 +99,11 @@ class OnnxConfig(ABC):
                 "end_logits": {0: "batch", 1: "sequence"},
             }
         ),
+        "semantic-segmentation": OrderedDict(
+            {
+                "logits": {0: "batch", 1: "sequence", 2: "height", 3: "width"},
+            }
+        ),
         "seq2seq-lm": OrderedDict({"logits": {0: "batch", 1: "decoder_sequence"}}),
         "sequence-classification": OrderedDict({"logits": {0: "batch"}}),
         "token-classification": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),