Update davit.py

3 years ago · 1327a6af4c
parent 7dd25eb402
commit 1327a6af4c
1 changed files with 2 additions and 2 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -508,7 +508,7 @@ class SpatialBlock(nn.Module):
        x = x.view(B, H * W, C)
        x = shortcut + self.drop_path(x)

-        x = self.cpe2(x)
+        x = self.cpe2(x.transpose(1, 2).view(B, C, H, W)).flatten(2).transpose(1, 2)
        if self.ffn:
            x = x + self.drop_path(self.mlp(self.norm2(x)))