Update davit.py

timm/models/davit.py

@@ -411,8 +411,10 @@ class DaViT(nn.Module):
         for stage_id, stage_param in enumerate(self.architecture):
             layer_offset_id = len(list(itertools.chain(*self.architecture[:stage_id])))
 
-            stage = nn.Sequential([
+            stage = nn.Sequential(
-                nn.Sequential([
+                nn.ModuleList([
+                    nn.Sequential(
+                        nn.ModuleList([
                             ChannelBlock(
                                 dim=self.embed_dims[item],
                                 num_heads=self.num_heads[item],
@@ -434,9 +436,10 @@ class DaViT(nn.Module):
                                 cpe_act=cpe_act,
                                 window_size=window_size,
                             ) if attention_type == 'spatial' else None
-                    for attention_id, attention_type in enumerate(attention_types)]
+                        for attention_id, attention_type in enumerate(attention_types)])
                     ) for layer_id, item in enumerate(stage_param)
                 ])
+            )
             
             self.stages.add_module(f'stage_{stage_id}', stage)