Update davit.py

3 years ago · 2faeaebfee
parent e7c5ab9d1e
commit 2faeaebfee
1 changed files with 1 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -607,7 +607,7 @@ class DaViT(nn.Module):
        outs = []
        for i, out in enumerate(x):
            H, W = sizes[i]
-            outs.append(x.view(-1, H, W, self.embed_dims[i]).permute(0, 3, 1, 2).contiguous())
+            outs.append(out.view(-1, H, W, self.embed_dims[i]).permute(0, 3, 1, 2).contiguous())
        
        return outs