Update ImageList to work with 3d tensors (facebookresearch#543)

Lyears · Mar 7, 2019 · 1a744b9 · 1a744b9
1 parent 43345f4
commit 1a744b9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/maskrcnn_benchmark/structures/image_list.py b/maskrcnn_benchmark/structures/image_list.py
@@ -41,6 +41,8 @@ def to_image_list(tensors, size_divisible=0):
         return tensors
     elif isinstance(tensors, torch.Tensor):
         # single tensor shape can be inferred
+        if tensors.dim() == 3:
+            tensors = tensors[None]
         assert tensors.dim() == 4
         image_sizes = [tensor.shape[-2:] for tensor in tensors]
         return ImageList(tensors, image_sizes)