Having a bit of information on the function might help a little bit with giving me an idea as to what you’re differentiating.
I have done calculating the 2nd derivative of a determinant function (which is why I asked if it’s that). When I made a custom version of that function it also suffered from near 0 values. Which I solved by using a custom torch.autograd.Function and manually defining the derivatives (with some torch.linalg.svd as well in order to handle singular matrices).
Could this be a potential avenue for you?