Trying to retain patch location info

So I have the following code:

Which basically creates patches of a 2D matrix of 3x3, then flattens the 3x3 patches to 9 elements.

image = torch.randn(1, 1, 5, 5)

kh, kw = 3, 3 # kernel size
dh, dw = 1, 1 # stride
patches = image.unfold(2, kh, dh).unfold(3, kw, dw)
patches = patches.contiguous().view(-1, 9)

I have been struggling to figure out how to keep the patch original position info along with the data on the patch. I have considered just trying to append a tensor of the positions to the tensor of the patches, however I can’t seem to figure out how to actually make that one work in code though. Basically, my goal is to give my network the patch location info along with the patches, but not sure what would be an effective way to do this. Thoughts?

Thanks,

Bob