Use F.fold to get the image as the patches just like for convolutional layers (it will be of shape batch x number of patches x size of a patch, then you implement your operation, and F.unfold to get back. The documentation (either for the functional or the modular interface) has an example for convolution.
Best regards
Thomas