Update davit.py

3 years ago · 55cfdd7b0d
parent 4851c92f48
commit 55cfdd7b0d
1 changed files with 1 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -226,7 +226,7 @@ class PatchEmbed(nn.Module):
        x = self.proj(x)
        #x = x.flatten(2).transpose(1, 2)
        if self.norm.normalized_shape[0] == self.embed_dim:
-            x = self.norm(x)
+            x = self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
        return x