From df4dc75dda3358fc4c6f7c1aa289b67c77ec0354 Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Tue, 25 Nov 2025 13:53:09 -0800 Subject: [PATCH] Add files via upload --- sam3/sam/transformer.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/sam3/sam/transformer.py b/sam3/sam/transformer.py index 3e96c28..dc5bf2d 100644 --- a/sam3/sam/transformer.py +++ b/sam3/sam/transformer.py @@ -282,9 +282,8 @@ class RoPEAttention(Attention): 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 @@ class RoPEAttention(Attention): # 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