Tweak over protective tests

torchmd · Oct 10, 2023 · f837503 · f837503
1 parent 0f0dbe0
commit f837503
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/tests/test_cfconv.py b/tests/test_cfconv.py
@@ -77,5 +77,5 @@ def test_cfconv(device, num_atoms, num_filters, num_rbfs, cutoff_upper):
     total.backward()
     grad = pos.grad.clone()
 
-    assert pt.allclose(ref_output, output, atol=5e-7)
-    assert pt.allclose(ref_grad, grad, atol=5e-7)
+    pt.testing.assert_close(ref_output, output, atol=1e-6, rtol=1e-5)
+    pt.testing.assert_close(ref_grad, grad, atol=1e-6, rtol=1e-5)
diff --git a/tests/test_optimize.py b/tests/test_optimize.py
@@ -52,5 +52,5 @@ def test_gn(device, num_atoms):
     # Execute the optimize model
     energy, gradient = model(elements, positions)
 
-    assert pt.allclose(ref_energy, energy, atol=5e-7)
-    assert pt.allclose(ref_gradient, gradient, atol=1e-5)
+    pt.testing.assert_close(ref_energy, energy, rtol=1e-5, atol=1e-5)
+    pt.testing.assert_close(ref_gradient, gradient, rtol=1e-4, atol=1e-5)