Skip to content

Commit df3484a

Browse files
authored
Merge pull request #166 from jhwei/main
Fix window_reverse to support batch inference when using onnx or tensorrt
2 parents af49aed + 45ec226 commit df3484a

1 file changed

Lines changed: 3 additions & 3 deletions

File tree

models/backbones/swin_v1.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -64,9 +64,9 @@ def window_reverse(windows, window_size, H, W):
6464
Returns:
6565
x: (B, H, W, C)
6666
"""
67-
B = int(windows.shape[0] / (H * W / window_size / window_size))
68-
x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
69-
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
67+
C = int(windows.shape[-1])
68+
x = windows.view(-1, H // window_size, W // window_size, window_size, window_size, C)
69+
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, H, W, C)
7070
return x
7171

7272

0 commit comments

Comments
 (0)