Update on "Don't use NonVariableTypeMode in custom ops"

Potentially fixes #37306

Differential Revision: [D21261946](https://our.internmc.facebook.com/intern/diff/D21261946/)

[ghstack-poisoned]

test/test_cpp_extensions_jit.py

@@ -13,6 +13,7 @@
 import torch.backends.cudnn
 import torch.utils.cpp_extension
 from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
+from torch.autograd.gradcheck import gradcheck
 
 
 TEST_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
@@ -832,7 +833,7 @@ def test_custom_compound_op_autograd(self):
         a = torch.randn(5, 5, requires_grad=True)
         b = torch.randn(5, 5, requires_grad=True)
 
-        gradcheck(torch.ops.my.add, [a, b])
+        gradcheck(torch.ops.my.add, [a, b], eps=1e-2)
 
 
 if __name__ == "__main__":