Update cvt.py

Author: fffffgggg54

timm/models/cvt.py

@@ -229,7 +229,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 = x.flatten(2).transpose(1, 2) + self.drop_path1(self.ls1(self.fw(attn(self.norm1(x)))))
+        x = x.flatten(2).transpose(1, 2) + self.drop_path1(self.ls1(self.fw_attn(self.norm1(x))))
         x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
 
         if self.use_cls_token:
@@ -395,4 +395,6 @@ def __init__(
                 use_cls_token = use_cls_token[stage_idx],
             )
             in_chs = dim
-            stages.append(stage)
+            stages.append(stage)
+        self.stages = nn.ModuleList(stages)
+