You could use a view for it:
input_windows_sorted = input_windows_sorted.contiguous().view(
    input_windows_sorted.size(0), -1, *input_windows_sorted.size()[-2:])

If the batch_size is known before, you can use it instead of input_windows_sorted.size(0), which makes the code a bit more readable.