From d06f5e86b601cffc37f9ae2e5065ac0dfd1cee21 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Mon, 2 Dec 2024 10:05:06 +0800 Subject: [PATCH] fix torchscript complaint --- timm/models/swin_transformer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/models/swin_transformer.py b/timm/models/swin_transformer.py index c6e0e119c..67c44aa11 100644 --- 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]