From c34d5cbd80483c480d8407420d99acf0f4ed6f11 Mon Sep 17 00:00:00 2001 From: Eric Jinks Date: Mon, 17 Feb 2020 11:47:19 +1000 Subject: [PATCH] Update segmentation_mask.py torch.bool --- maskrcnn_benchmark/structures/segmentation_mask.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/maskrcnn_benchmark/structures/segmentation_mask.py b/maskrcnn_benchmark/structures/segmentation_mask.py index ba1290b91..a37473635 100644 --- 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()