Say you have a gray image tensor of shape (1, 1, 128, 128) . What I would like to do here is to sample in each h, w dimension with stride=2, which would then make 4 sub-images of size (1, 1, 64, 64) depending on where the indexing starts. Then, I would like to batch them to finally form the tensor of size (4, 1, 64, 64). For now, I’m using double for loops which is presumably inefficient. Is there an efficient way to vectorize this using functions implemented in pytorch?
The function I’m using right now:
def img_subsample(img, ratio=2):
    '''
    ex) Subsamples img of size 128 * 128,
    output is batched into 4 * 1 * 64 * 64
    Args:
        img: input of size ex) 128 * 128
        ratio(int): Subsampling ratio
    Returns: batched, subsampled img
    '''
    sz = img.shape[-1]
    batched_output = torch.zeros([ratio**2, 1, (sz//ratio), (sz//ratio)])
    for i in range(ratio):
        for j in range(ratio):
            batched_output[ratio*i+j, :, :, :] = img[i::ratio, j::ratio]
    return batched_output