Sorry I realized that the code is not working exactly as expected, can you please let me know your opinion.
So here i provide an example:
for example:
kh, kw = 2, 2
dh, dw = 2, 2
input = torch.rand(1,2,4,4)
input_windows = input.unfold(2, kh, dh).unfold(3, kw, dw)
input_windows = input_windows.contiguous().view(*input_windows.size()[:-2], -1)
input_windows_sorted = input_windows.sort(descending=True)[0]
input_windows_sorted = input_windows_sorted.view(*input.size())
input_windows_sorted = input_windows_sorted.transpose(2, 3)
input_windows_sorted = input_windows_sorted.unfold(3, kh, kw)
print(input_windows_sorted.size())

Output: torch.Size([1, 2, 4, 2, 2])

but when i change the size of channels i see:
input = torch.rand(1,2,6,6)
input_windows = input.unfold(2, kh, dh).unfold(3, kw, dw)
input_windows = input_windows.contiguous().view(*input_windows.size()[:-2], -1)
input_windows_sorted = input_windows.sort(descending=True)[0]
input_windows_sorted = input_windows_sorted.view(*input.size())
input_windows_sorted = input_windows_sorted.transpose(2, 3)
input_windows_sorted = input_windows_sorted.unfold(3, kh, kw)
print(input_windows_sorted.size())

OutPut: torch.Size([1, 2, 6, 3, 2])

but the output should be
OutPut: torch.Size([1, 2, 4, 3, 3)