Wouldn’t something like this do the trick:
class CustomConv(torch.nn.Module):
    def __init__(p, *args, **kwargs):
        super().__init__(self)
        self.conv = torch.nn.Conv2d(*args, **kwargs)
        self.p = p

    def forward(self, x, z):
        diff = x-z
        p_norm = torch.abs(x-z).pow(self.p)
        return self.conv(p_norm) 

There might be some typos or little mistakes as I’m typing from my smartphone.