Update davit.py

timm/models/davit.py

@@ -35,11 +35,18 @@ from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 __all__ = ['DaViT']
 
 class SequentialWithSize(nn.Sequential):
+
+    def __init__(*args, **kwargs):
+        super(MySequential, self).__init__(*args, **kwargs)
+
     def forward(self, x: Tensor, size: Tuple[int, int]):
         for module in self._modules.values():
+            x, size = module(x, size)
+            '''
             output = module(x, size)
             x : Tensor = output[0]
             size : Tuple[int, int] = output[1]
+            '''
         return x, size