Thanks for your reply! Here is my RigidityExp function:
class RigidityExp(torch.autograd.Function):
 
    @staticmethod
    def forward(ctx, input, eta, beta, cutoff):
        """
        Input:
            - input  : (D_in, N, M) matrix, which will be the output of sparse2tensor functionself.
            - eta    : (D_in, 1, 1) matrix.
            - beta   :  constant
            - cutoff :  constant
        Output:
            - output : A (1, D_in) matrix
        """
        ctx.save_for_backward(input, eta)
        ctx.beta =  beta
        ctx.cutoff = cutoff
        X = torch.exp(-(input/eta).pow(beta))
        M = torch.zeros_like(X)
        Y = torch.where(input < cutoff, X, M)
        output = torch.sum(torch.sum(Y, dim =-1), dim=-1).unsqueeze(0)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input, eta = ctx.saved_tensors
        exponent = (input/eta).pow(ctx.beta)
        X = torch.exp(-exponent)
        M = torch.zeros_like(X)
        M0 = torch.ones_like(X)
        Y = torch.where(input < ctx.cutoff, X, M)
        input0 = torch.where(input != 0, input, M0)
        Z = (ctx.beta * Y * exponent) / eta
        Z0 = -(ctx.beta * Y * exponent)/ input0

        matrix_eta = torch.sum(torch.sum(Z, dim =-1), dim=-1).unsqueeze(0)
        matrix_eta0 = torch.sum(torch.sum(Z0, dim =-1), dim=-1).unsqueeze(0)

        grad_input = (grad_output.unsqueeze(0).permute(2,0,1) * matrix_eta0)
        grad_eta = (grad_output * matrix_eta).unsqueeze(0).permute(2,0,1)
        return grad_input, grad_eta, None, None

from torch.autograd import gradcheck
rigidity = RigidityExp.apply
a = ((15 - 2.5) * torch.rand(1, 36, device = device,requires_grad=True) + 2.5).double().unsqueeze(0).permute(2, 1, 0)
input = (sparse2tensor(X_test_Distance[0,:].reshape(1,-1)).double(), a, 2.5, 12)
test = gradcheck(rigidity, input, eps=1e-3, atol=1e-3)
print(test)