Update davit.py

2 years ago · a5e4114cf5
parent fa45868239
commit a5e4114cf5
1 changed files with 1 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -218,7 +218,7 @@ def window_reverse(windows : Tensor, window_size: int, H: int, W: int):
    
    B : float = (windows.size(dim=0) / (H * W / window_size / window_size))
    x = windows.view(int(B), 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)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(int(B), H, W, -1)
    return x