From f5e8f8246516c5245231acab5d5284cc81992fbe Mon Sep 17 00:00:00 2001
From: Fredo Guan <fredo.guan@hotmail.com>
Date: Thu, 8 Dec 2022 17:41:25 -0800
Subject: [PATCH] Update davit.py

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

diff --git a/timm/models/davit.py b/timm/models/davit.py
index 57498f4d..dab88d32 100644
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@@ -216,8 +216,14 @@ def window_reverse(windows : Tensor, window_size: int, H: int, W: int):
         x: (B, H, W, C)
     """
     
-    B = torch.floor(torch.tensor(windows.size(dim=0)) / (H * W / window_size / window_size)).int().item()
-    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    #B = torch.floor(torch.tensor(windows.size(dim=0)) / (H * W / window_size / window_size)).int().item()
+    x = windows.view(
+        torch.floor(torch.tensor(windows.size(dim=0)) / (H * W / window_size / window_size)).int().item(),
+        H // window_size,
+        W // window_size,
+        window_size,
+        window_size,
+        -1)
     x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
     return x