fter partition with [B * num_windows, window_size, window_size, C].
        (Hp, Wp): padded height and width before partition
    r