As I mentioned on Twitter, I’d start with something based on numpy and adapt/optimize for pytorch. Something like https://github.com/jabooth/patchmatch/blob/master/patchmatch.py for example.
It will be definitely slower than using hand-made kernels, but will definitely be easier to debug.
Good luck!