Merge pull request #4 from Jinksi/segm-index-fix

Segmentation index error fix
Jinksi · Feb 20, 2020 · fcc5c68 · fcc5c68
2 parents fa586d4 + c34d5cb
commit fcc5c68
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/maskrcnn_benchmark/structures/segmentation_mask.py b/maskrcnn_benchmark/structures/segmentation_mask.py
@@ -195,7 +195,9 @@ def __getitem__(self, item):
         else:
             # advanced indexing on a single dimension
             selected_polygons = []
-            if isinstance(item, torch.Tensor) and item.dtype == torch.uint8:
+            if isinstance(item, torch.Tensor) and (
+                item.dtype == torch.uint8 or item.dtype == torch.bool
+            ):
                 item = item.nonzero()
                 item = item.squeeze(1) if item.numel() > 0 else item
                 item = item.tolist()