And for auto.grad part if use .double() for input, then I will get True when I check grad. And I will get an error if I do not double input and param.
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)
>>True

And
from torch.autograd import gradcheck
Rigidity = RigidityExp.apply
a = ((15 - 2.5) * torch.rand(1, 36,device = device,requires_grad=True) + 2.5).unsqueeze(0).permute(2, 1, 0)
input = (sparse2tensor(X_test_Distance[0,:].reshape(1,-1)), a, 2.5, 12)
test = gradcheck(Rigidity, input, eps=1e-3, atol=1e-3)
print(test)
>>/home/name/anaconda3/lib/python3.6/site-packages/torch/autograd/gradcheck.py:170: UserWarning: At least one of the inputs that requires gradient is not of double precision floating point. This check will likely fail if all the inputs are not of double precision floating point. 
  'At least one of the inputs that requires gradient '
Traceback (most recent call last):
  File "check_grad.py", line 229, in <module>
    test = gradcheck(linear, input, eps=1e-3, atol=1e-3)
  File "/home/name/anaconda3/lib/python3.6/site-packages/torch/autograd/gradcheck.py", line 205, in gradcheck
    'numerical:%s\nanalytical:%s\n' % (i, j, n, a))
  File "/home/name/anaconda3/lib/python3.6/site-packages/torch/autograd/gradcheck.py", line 185, in fail_test
    raise RuntimeError(msg)
RuntimeError: Jacobian mismatch for output 0 with respect to input 0,
numerical:tensor([[0.2441, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.7324, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.2441,  ..., 0.0000, 0.0000, 0.0000],
        ...,
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000]])
analytical:tensor([[4.5538e-01, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00, 0.0000e+00,
         0.0000e+00],
        [0.0000e+00, 7.2516e-01, 0.0000e+00,  ..., 0.0000e+00, 0.0000e+00,
         0.0000e+00],
        [0.0000e+00, 0.0000e+00, 3.0459e-01,  ..., 0.0000e+00, 0.0000e+00,
         0.0000e+00],
        ...,
        [0.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00, 0.0000e+00,
         0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00, 1.8696e-04,
         0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00, 0.0000e+00,
         0.0000e+00]])