Update davit.py

3 years ago · d5b74ae34f
parent 9d20d6ef6a
commit d5b74ae34f
1 changed files with 1 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -419,7 +419,7 @@ class DaViTStage(nn.Module):
                        cpe_act=cpe_act
                    ))
            
-            stage_blocks.append(nn.ModuleList(dual_attention_block))
+            stage_blocks.append(SequentialWithSize(*dual_attention_block))
            
        self.blocks = SequentialWithSize(*stage_blocks)