For the newer versions of pytorch, any torch.autograd.Function needs to be implemented with static member functions. Here is a pretty comprehensive example, n.b. the extra custom_[f|b]wd decorators are for amp and won’t be required if you aren’t using it.
class CorrelationFunction(torch.autograd.Function):
    """
    Typical Parameters: pad_size=3, kernel_size=3, max_displacement=20,
    stride1=1, stride2=2, corr_multiply=1
    """

    @staticmethod
    @torch.cuda.amp.custom_fwd
    def forward(ctx, input1, input2, pad_size=3, kernel_size=3,
                max_displacement=20, stride1=1, stride2=2, corr_multiply=1):
        ctx.save_for_backward(input1, input2)

        ctx.pad_size = pad_size
        ctx.kernel_size = kernel_size
        ctx.max_displacement = max_displacement
        ctx.stride1 = stride1
        ctx.stride2 = stride2
        ctx.corr_multiply = corr_multiply

        output = torch.ops.cerberus.correlation(
            input1, input2, pad_size, kernel_size,
            max_displacement, stride1, stride2, corr_multiply)

        return output

    @staticmethod
    @torch.cuda.amp.custom_bwd
    def backward(ctx, grad_outputs):
        input1, input2 = ctx.saved_tensors

        grad_input1, grad_input2 = torch.ops.cerberus.correlation_backward(
            input1, input2, grad_outputs, ctx.pad_size, ctx.kernel_size,
            ctx.max_displacement, ctx.stride1, ctx.stride2, ctx.corr_multiply)

        return grad_input1, grad_input2, None, None, None, None, None, None

class Correlation(torch.nn.Module):
    def __init__(self, pad_size=0, kernel_size=0, max_displacement=0,
                 stride1=1, stride2=2, corr_multiply=1):
        super(Correlation, self).__init__()

        self.pad_size = pad_size
        self.kernel_size = kernel_size
        self.max_displacement = max_displacement
        self.stride1 = stride1
        self.stride2 = stride2
        self.corr_multiply = corr_multiply

    def forward(self, input1, input2):
        return CorrelationFunction.apply(
            input1, input2, self.pad_size, self.kernel_size,
            self.max_displacement, self.stride1, self.stride2, self.corr_multiply)