Update davit.py

timm/models/davit.py

@@ -410,12 +410,12 @@ class DaViT(nn.Module):
             for i in range(self.num_stages)])
 
         #main_blocks = []
-        self.main_blocks = nn.ModuleList()
+        self.stages = nn.Sequential()
-        for block_id, block_param in enumerate(self.architecture):
+        for stage_id, stage_param in enumerate(self.architecture):
-            layer_offset_id = len(list(itertools.chain(*self.architecture[:block_id])))
+            layer_offset_id = len(list(itertools.chain(*self.architecture[:stage_id])))
 
-            block = nn.ModuleList([
+            stage = nn.Sequential([
-                nn.ModuleList([
+                nn.Sequential([
                     ChannelBlock(
                         dim=self.embed_dims[item],
                         num_heads=self.num_heads[item],
@@ -438,12 +438,12 @@ class DaViT(nn.Module):
                         window_size=window_size,
                     ) if attention_type == 'spatial' else None
                     for attention_id, attention_type in enumerate(attention_types)]
-                ) for layer_id, item in enumerate(block_param)
+                ) for layer_id, item in enumerate(stage_param)
             ])
             
-            self.main_blocks.add_module(f'block', block)
+            self.main_blocks.append(stage)
             
-            self.feature_info += [dict(num_ch=self.embed_dims[block_id], reduction = 2, module=f'block.{block_id}')]
+            self.feature_info += [dict(num_ch=self.embed_dims[stage_id], reduction = 2, module=f'stage.{stage_id}')]
         #self.main_blocks = nn.ModuleList(main_blocks)
         
         '''