Update davit.py

timm/models/davit.py

@@ -411,10 +411,8 @@ 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.ModuleList([
+                nn.Sequential([
-                    nn.Sequential(
-                        nn.ModuleList([
                     ChannelBlock(
                         dim=self.embed_dims[item],
                         num_heads=self.num_heads[item],
@@ -436,10 +434,9 @@ 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)