You can slightly change your code to:
import torch
from nnd import NNDModule
data1 = torch.rand((2, 3, 5), requires_grad=True, dtype=torch.double, device="cuda") #[b, c, n]
data2 = torch.rand((2, 3, 4), requires_grad=True, dtype=torch.double, device="cuda") #[b, c, m]
dist = NNDModule()
dist.double()
print(torch.autograd.gradcheck(dist, (data1, data2)))

You shouldn’t need to change the eps.

but the torch.max() can also pass the gradcheck

It does as long as you’re far-enough from a point where the max value changes. If the finite difference actually hit a point where you get a different max, then the test will fail.

My function can only receive float precision input

You actually give double types above no? Do you convert them back to float before doing the actual computation?

The formula as the backward pass:

The forward pass actually does an outer product between x and y right? I am surprised not to see one here.