Update davit.py

3 years ago · f47f6fce31
parent dd6531c525
commit f47f6fce31
1 changed files with 2 additions and 2 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -411,8 +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.ModuleList([
-                nn.ModuleList([
+            stage = nn.Sequential([
+                nn.Sequential([
                    ChannelBlock(
                        dim=self.embed_dims[item],
                        num_heads=self.num_heads[item],