fix ut

deepmodeling · Feb 1, 2024 · 9dd10c9 · 9dd10c9
1 parent 600349b
commit 9dd10c9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/source/tests/pt/test_rotation.py b/source/tests/pt/test_rotation.py
@@ -111,7 +111,7 @@ def test_rotation(self):
         result1 = self.model(**get_data(self.origin_batch))
         result2 = self.model(**get_data(self.rotated_batch))
         rotation = torch.from_numpy(self.rotation).to(env.DEVICE)
-        self.assertAlmostEqual(result1["energy"], result2["energy"])
+        torch.testing.assert_close(result1["energy"], result2["energy"])
         if "force" in result1:
             torch.testing.assert_close(
                 result2["force"][0], torch.matmul(rotation, result1["force"][0].T).T