draw_segmentation_masks: Add support for GPU

[mshr-h]

🚀 The feature

Adding support for draw_segmentation_masks computation on GPU by creating the colors tensor in the same device as the input image tensor.

Motivation, pitch

The current draw_segmentation_masks function requires all the input tensors to be located in the CPU because the colors tensor is created in the CPU.
By placing the colors tensor in the same device as the input image tensor, we could run all the computation in GPU if the input tensors are in the GPU.

-    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)]

Alternatives

No response

Additional context

No response