How to perform shift operation from ShiftNet image inpainting paper

I am facing difficulty in implementing a shift layer from shiftnet image inpainting paper. the link for the paper is shiftnet image inpainting paper.
I want to perform cross-correlation operation in the third last layer of the U-Net type architecture given in this paper.

could someone please help me figure out how can I performs this shift operation in PyTorch?
Thanks in advance

The shift vector is computed in the paper as given below

the idea is basically to perform the nearest neighbour search for each pixel in the corrupted region and then replace the nearest neighbour element inplace and further concatenating with the already upsampled and skip connected feature map.
But I am unable to figure out a way to do so in pytorch!