add support for 4D output (open-mmlab#150)

bowenroom · Sep 23, 2020 · 588a2c0 · 588a2c0
1 parent 7baed65
commit 588a2c0
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