fix torchscript complaint

huggingface · Dec 2, 2024 · d06f5e8 · d06f5e8
1 parent 7fd7206
commit d06f5e8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/timm/models/swin_transformer.py b/timm/models/swin_transformer.py
@@ -291,7 +291,7 @@ def get_attn_mask(self, x: Optional[torch.Tensor] = None) -> Optional[torch.Tens
                 dtype = x.dtype
             else:
                 H, W = self.input_resolution
-                device = "cpu"
+                device = torch.device("cpu")
                 dtype = None
             H = math.ceil(H / self.window_size[0]) * self.window_size[0]
             W = math.ceil(W / self.window_size[1]) * self.window_size[1]