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

Author: drcut

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