I think I got this  I moved the channel axis to the front s.t. it would be safe and it passes my manual tests.
def img_to_batch(img, patch_size: int):
    _, ch, height, width = img.shape
    assert height%patch_size == 0 and width%patch_size == 0, (
        'img_to_batch: dims must be dividable by patch_size. {}%{}!=0'.format(
            img.shape, patch_size))
    assert img.dim() == 4
    return img.unfold(2, patch_size, patch_size).unfold(
        3, patch_size, patch_size).transpose(1,0).reshape(
            ch, -1, patch_size, patch_size).transpose(1,0)