From 08f30ea497cfe7279ce8c1310620abdc31a20db4 Mon Sep 17 00:00:00 2001 From: robin Han Date: Wed, 23 Sep 2020 17:01:20 +0800 Subject: [PATCH] add support for 4D output (#150) --- mmseg/models/segmentors/encoder_decoder.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mmseg/models/segmentors/encoder_decoder.py b/mmseg/models/segmentors/encoder_decoder.py index 9adf65bd0..2284906e3 100644 --- a/mmseg/models/segmentors/encoder_decoder.py +++ b/mmseg/models/segmentors/encoder_decoder.py @@ -265,6 +265,8 @@ class EncoderDecoder(BaseSegmentor): 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