Update davit.py

3 years ago · 1fa1e01646
parent f5e8f82465
commit 1fa1e01646
1 changed files with 5 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -224,7 +224,11 @@ def window_reverse(windows : Tensor, window_size: int, H: int, W: int):
        window_size,
        window_size,
        -1)
-    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(
+        torch.floor(torch.tensor(windows.size(dim=0)) / (H * W / window_size / window_size)).int().item(),
+        H,
+        W,
+        -1)
    return x