made device=self.device

aim-uofa · Jun 16, 2022 · f60e856 · f60e856
1 parent c0b2092
commit f60e856
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/adet/modeling/condinst/condinst.py b/adet/modeling/condinst/condinst.py
@@ -304,7 +304,7 @@ def add_bitmasks_from_boxes(self, instances, images, image_masks, im_h, im_w):
             per_im_bitmasks = []
             per_im_bitmasks_full = []
             for per_box in per_im_boxes:
-                bitmask_full = torch.zeros((im_h, im_w)).to(self.device).float()
+                bitmask_full = torch.zeros((im_h, im_w), device=self.device).float()
                 bitmask_full[int(per_box[1]):int(per_box[3] + 1), int(per_box[0]):int(per_box[2] + 1)] = 1.0
 
                 bitmask = bitmask_full[start::stride, start::stride]