You could try to use PyTorch methods, if the numpy equivalents are available, or you would have to write a custom autograd.Function with the forward and backward pass manually.
Could you post the method so that we can have a look, which numpy ops you are using?