I have the following code in my nn.Module.
x = torch.cdist(a,b)
y = torch.softmax(x)
These two are differentiable, but due to the size of x, x/y need a lot of GPU memory, causing a OOM during backprop. Since I’m not sure if the checkpoint can help this (getting NaN with ddp) and x/y are very sparse (many near zero values), I was thinking of a custom operator in the following fashion, but cannot find a doc on how to call backward of existing pytorch constructs in a custom operator. Can anyone point me out some examples on how to complete this?
class MyFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, a, b):        
        x = torch.cdist(a,b)
        ctx.save_for_backward(a,b, sparse(x))
        return torch.softmax(x)

    @staticmethod
    def backward(ctx, g):
     ??