Custom nn.Conv2d

I think the easiest way would be to use unfold to create patches and apply your custom functions on each patch.
This post might be a good starter.