add CUDA support for draw_segmentation_masks

pytorch · Jun 20, 2023 · 926cbc3 · 926cbc3
1 parent b9a1984
commit 926cbc3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchvision/utils.py b/torchvision/utils.py
@@ -304,7 +304,7 @@ def draw_segmentation_masks(
         return image
 
     out_dtype = torch.uint8
-    colors = [torch.tensor(color, dtype=out_dtype) for color in _parse_colors(colors, num_objects=num_masks)]
+    colors = [torch.tensor(color, dtype=out_dtype, device=image.device) for color in _parse_colors(colors, num_objects=num_masks)]
 
     img_to_draw = image.detach().clone()
     # TODO: There might be a way to vectorize this