Update davit.py

3 years ago · 38ea27001f
parent d03e8cc94d
commit 38ea27001f
1 changed files with 1 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -421,7 +421,7 @@ 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([
+            stage = MySequential(*[
                MySequential(*[
                    ChannelBlock(
                        dim=self.embed_dims[item],