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
```