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

---
 timm/models/davit.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/timm/models/davit.py b/timm/models/davit.py
index 2e3e33ba..1028c82f 100644
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@@ -508,11 +508,11 @@ 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)))
             
-        x = x.transpose(1, 2).view(B, C, H, W)    
+        x = x.transpose(1, 2).view(B, C, H, W)
         
         return x