From 8f02ec1bc0806cbaa0788beab9fe7574411cde73 Mon Sep 17 00:00:00 2001 From: danikvh Date: Mon, 4 Aug 2025 18:48:11 +0200 Subject: [PATCH] Fix device mismatch in bounding box computation --- cellSAM/AnchorDETR/util/box_ops.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cellSAM/AnchorDETR/util/box_ops.py b/cellSAM/AnchorDETR/util/box_ops.py index 9ffe9a7..c62f143 100644 --- a/cellSAM/AnchorDETR/util/box_ops.py +++ b/cellSAM/AnchorDETR/util/box_ops.py @@ -81,8 +81,8 @@ def masks_to_boxes(masks): h, w = masks.shape[-2:] - y = torch.arange(0, h, dtype=torch.float) - x = torch.arange(0, w, dtype=torch.float) + y = torch.arange(0, h, dtype=torch.float, device=masks.device) + x = torch.arange(0, w, dtype=torch.float, device=masks.device) y, x = torch.meshgrid(y, x) x_mask = (masks * x.unsqueeze(0)) @@ -93,4 +93,4 @@ def masks_to_boxes(masks): y_max = y_mask.flatten(1).max(-1)[0] y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] - return torch.stack([x_min, y_min, x_max, y_max], 1) + return torch.stack([x_min, y_min, x_max, y_max], 1) \ No newline at end of file