Always some torchscript issues

timm/layers/space_to_depth.py

@@ -3,6 +3,8 @@
 
 
 class SpaceToDepth(nn.Module):
+    bs: torch.jit.Final[int]
+
     def __init__(self, block_size=4):
         super().__init__()
         assert block_size == 4
@@ -12,7 +14,7 @@ def forward(self, x):
         N, C, H, W = x.size()
         x = x.view(N, C, H // self.bs, self.bs, W // self.bs, self.bs)  # (N, C, H//bs, bs, W//bs, bs)
         x = x.permute(0, 3, 5, 1, 2, 4).contiguous()  # (N, bs, bs, C, H//bs, W//bs)
-        x = x.view(N, C * (self.bs ** 2), H // self.bs, W // self.bs)  # (N, C*bs^2, H//bs, W//bs)
+        x = x.view(N, C * self.bs * self.bs, H // self.bs, W // self.bs)  # (N, C*bs^2, H//bs, W//bs)
         return x
 
 

timm/models/pvt_v2.py

@@ -66,6 +66,8 @@ def forward(self, x, feat_size: List[int]):
 
 
 class Attention(nn.Module):
+    fused_attn: torch.jit.Final[bool]
+
     def __init__(
             self,
             dim,