oh xd i feel stupid

Author: fffffgggg54

timm/models/cvt.py

@@ -289,7 +289,7 @@ def fw_attn(self, x: torch.Tensor, cls_token: Optional[torch.Tensor]) -> torch.T
     def forward(self, x: torch.Tensor, cls_token: Optional[torch.Tensor]) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
         B, C, H, W = x.shape
 
-        x = torch.cat((cls_token, x.flatten(2).transpose(1, 2)), dim=1) if cls_token is not None else x.flatten(2).transpose(1, 2) \
+        x = (torch.cat((cls_token, x.flatten(2).transpose(1, 2)), dim=1) if cls_token is not None else x.flatten(2).transpose(1, 2)) \
             + self.drop_path1(self.ls1(self.fw_attn(self.norm1(x), cls_token)))
         x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))