From 148cffa23627d7476c1ddc83a43afe7bd0c3c584 Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Tue, 25 Nov 2025 13:52:43 -0800 Subject: [PATCH] Add files via upload --- sam3/model/decoder.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/sam3/model/decoder.py b/sam3/model/decoder.py index c8b1657..7e05304 100644 --- a/sam3/model/decoder.py +++ b/sam3/model/decoder.py @@ -277,10 +277,8 @@ class TransformerDecoder(nn.Module): if resolution is not None and stride is not None: feat_size = resolution // stride - coords_h, coords_w = self._get_coords( - feat_size, feat_size, device="cuda" - ) - self.compilable_cord_cache = (coords_h, coords_w) + # Build cache lazily on the actual reference_boxes device to avoid + # mixing CPU/GPU tensors when users run on CPU. self.compilable_stored_size = (feat_size, feat_size) self.roi_pooler = ( @@ -332,7 +330,11 @@ class TransformerDecoder(nn.Module): H, W = feat_size boxes_xyxy = box_cxcywh_to_xyxy(reference_boxes).transpose(0, 1) bs, num_queries, _ = boxes_xyxy.shape - if self.compilable_cord_cache is None: + if ( + self.compilable_cord_cache is None + or any(c.device != reference_boxes.device for c in self.compilable_cord_cache) + or self.compilable_stored_size != (H, W) + ): self.compilable_cord_cache = self._get_coords(H, W, reference_boxes.device) self.compilable_stored_size = (H, W) @@ -345,11 +347,12 @@ class TransformerDecoder(nn.Module): else: # cache miss, will create compilation issue # In case we're not compiling, we'll still rely on the dict-based cache - if feat_size not in self.coord_cache: - self.coord_cache[feat_size] = self._get_coords( + cache_key = (feat_size, reference_boxes.device) + if cache_key not in self.coord_cache: + self.coord_cache[cache_key] = self._get_coords( H, W, reference_boxes.device ) - coords_h, coords_w = self.coord_cache[feat_size] + coords_h, coords_w = self.coord_cache[cache_key] assert coords_h.shape == (H,) assert coords_w.shape == (W,)