set bs to 1 if inference (#1594)

dmlc · Jan 23, 2021 · 426a8f9 · 426a8f9
1 parent ecf4916
commit 426a8f9
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/gluoncv/model_zoo/rcnn/mask_rcnn/mask_rcnn.py b/gluoncv/model_zoo/rcnn/mask_rcnn/mask_rcnn.py
@@ -82,7 +82,11 @@ def hybrid_forward(self, F, x):
         # (B * N, C, MS, MS)
         x = self.mask(x)
         # (B * N, C, MS, MS) -> (B, N, C, MS, MS)
-        x = x.reshape((-4, self._batch_images, -1, 0, 0, 0))
+        if autograd.is_training():
+            x = x.reshape((-4, self._batch_images, -1, 0, 0, 0))
+        else:
+            # always use batch_size = 1 for inference
+            x = x.reshape((-4, 1, -1, 0, 0, 0))
         return x
 
     def reset_class(self, classes, reuse_weights=None):