Hi,
I am not quite sure if I understood your problem correctly, but it seems the torch.nn.Unfold and Fold function is what you’re looking for. If you have the problem that overlapping tiles are added together you can create a mask and divide by it to obtain the original input.
For example:
img     = torch.ones(1, 3,1024,1024)
mask    = torch.ones_like(img) ## Create mask to correct for areas which are going to be added 
# use torch.nn.Unfold
p_size  = 128
stride  = p_size//2
unfold  = nn.Unfold(kernel_size=(p_size, p_size), stride=stride)
# Apply to mask and original image
mask_p  = unfold(mask)
patches = unfold(img)
patches.shape

outputs shape of torch.Size([1, 49152, 225]) - for using the tiles individually you’ve to reshape them:
patches.reshape(3, p_size, p_size, -1).permute(3, 0, 1, 2)

to get a shape of torch.Size([225, 3, 128, 128]). Merging the patches/tiles back to the original image:
fold     = nn.Fold(output_size=(img.shape[2], img.shape[3]), kernel_size=(p_size, p_size), stride=stride)
img_back = fold(patches)/fold(mask_p)

If you don’t divide by the mask you will end up with areas where overlapping patches/tiles are added together. It isn’t the fanciest way but it works reasonable for different strides and patch sizes.