diff --git a/models/encoder/swin_transformer_v2.py b/models/encoder/swin_transformer_v2.py
index 79b3610..0a663f2 100644
--- a/models/encoder/swin_transformer_v2.py
+++ b/models/encoder/swin_transformer_v2.py
@@ -148,7 +148,8 @@ def forward(self, x, mask=None):
 
         # cosine attention
         attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))
-        logit_scale = torch.clamp(self.logit_scale, max=torch.log(torch.tensor(1. / 0.01))).exp()
+        max_val = torch.log(torch.tensor(1. / 0.01, device=self.logit_scale.device))
+        logit_scale = torch.clamp(self.logit_scale, max=max_val).exp()
         attn = attn * logit_scale
 
         relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads)
@@ -632,4 +633,4 @@ def build_swinv2_encoder(depths, embed_dim, num_heads, drop_path_rate, pretraine
         frozen_stages=-1,
         out_features=["stage1", "stage2", "stage3", "stage4", "stage5"],
         use_checkpoint=use_checkpoint,
-    )
\ No newline at end of file
+    )