You could try out your method by creating each window patch via unfold, apply your method, and reshape it to the output shape.
This post gives you an example of this approach.
Note that this approach will most likely be slow, but might be good enough to experiment with your method. If you need a more performance you might want to write a custom extension.