Thank you @Nikronic !
If I understand this correctly, then the img_to_batch function I wrote above would be simply:
def img_to_batch(imgtensor, patch_size: int):
    _, ch, height, width = imgtensor.shape
    assert height%patch_size == 0 and width%patch_size == 0, (
        'img_to_batch: dims must be dividable by patch_size. {}%{}!=0'.format(
            imgtensor.shape, patch_size))
    return imgtensor.unfold(2,patch_size,patch_size).unfold(
        3,patch_size,patch_size).contiguous().view(
            ch,-1,patch_size,patch_size).permute((1,0,2,3))       

Note that I made a slight modification from the equivalent of “x.unfold(2, 2, 2)[0].unfold(2,” to “x.unfold(2, 2, 2).unfold(3,”, I think that it can then also handle arbitrary batch size (and just like the channels they wouldn’t get mixed up). Correct me if I’m wrong.