Update davit.py

3 years ago · 619b207bc2
parent c579c62c32
commit 619b207bc2
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.Sequential([
-                nn.Sequential([
+            stage = nn.Sequential(*[
+                nn.Sequential(*[
                    ChannelBlock(
                        dim=self.embed_dims[item],
                        num_heads=self.num_heads[item],