Add files via upload

Author: 1038lab

sam3/sam/transformer.py

@@ -282,9 +282,8 @@ def __init__(
         self.compute_cis = partial(
             compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta
         )
-        device = torch.device("cuda") if torch.cuda.is_available() else None
         self.freqs_cis = self.compute_cis(
-            end_x=feat_sizes[0], end_y=feat_sizes[1], device=device
+            end_x=feat_sizes[0], end_y=feat_sizes[1], device=None
         )
         if self.use_rope_real:
             self.freqs_cis_real = self.freqs_cis.real
@@ -306,10 +305,14 @@ def forward(
 
         # Apply rotary position encoding
         w = h = math.sqrt(q.shape[-2])
-        if self.freqs_cis.shape[0] != q.shape[-2]:
+        if (
+            self.freqs_cis.device != q.device
+            or self.freqs_cis.shape[0] != q.shape[-2]
+        ):
             self.freqs_cis = self.compute_cis(end_x=w, end_y=h, device=q.device)
-            self.freqs_cis_real = self.freqs_cis.real
-            self.freqs_cis_imag = self.freqs_cis.imag
+            if self.use_rope_real:
+                self.freqs_cis_real = self.freqs_cis.real
+                self.freqs_cis_imag = self.freqs_cis.imag
         if q.shape[-2] != k.shape[-2]:
             assert self.rope_k_repeat