From 4851c92f48779ae49cf2c515b5520344f15f3238 Mon Sep 17 00:00:00 2001
From: Fredo Guan <fredo.guan@hotmail.com>
Date: Sat, 10 Dec 2022 20:48:50 -0800
Subject: [PATCH] Update davit.py

---
 timm/models/davit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/timm/models/davit.py b/timm/models/davit.py
index 1636bfd3..d764b6e6 100644
--- 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.flatten(2).transpose(1, 2)).transpose(1, 2).view(B, C, H, W)
+            x = self.norm(x)
         return x