add support for 4D output (#150)

open-mmlab · Sep 23, 2020 · 08f30ea · 08f30ea
1 parent a2738fd
commit 08f30ea
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/mmseg/models/segmentors/encoder_decoder.py b/mmseg/models/segmentors/encoder_decoder.py
@@ -265,6 +265,8 @@ def simple_test(self, img, img_meta, rescale=True):
         seg_logit = self.inference(img, img_meta, rescale)
         seg_pred = seg_logit.argmax(dim=1)
         if torch.onnx.is_in_onnx_export():
+            # our inference backend only support 4D output
+            seg_pred = seg_pred.unsqueeze(0)
             return seg_pred
         seg_pred = seg_pred.cpu().numpy()
         # unravel batch dim