How do I create / fix my custom 4 dimensional reflection padding function for PyTorch NCHW tensors?

Question

In order to apply reflection padding to all 4 dimensions of a tensor, I put together a custom padding function. However, it doesn\'t seem to handle padding the batch dimensi