add stress

ACEsuit · Nov 22, 2024 · 214ac8f · 214ac8f
1 parent 4ea5f9b
commit 214ac8f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/test_cueq.py b/tests/test_cueq.py
@@ -120,6 +120,8 @@ def test_bidirectional_conversion(
         torch.testing.assert_close(out_cueq["energy"], out_e3nn_back["energy"])
         torch.testing.assert_close(out_e3nn["forces"], out_cueq["forces"])
         torch.testing.assert_close(out_cueq["forces"], out_e3nn_back["forces"])
+        torch.testing.assert_close(out_e3nn["stress"], out_cueq["stress"])
+        torch.testing.assert_close(out_cueq["stress"], out_e3nn_back["stress"])
 
         # Test backward pass equivalence
         loss_e3nn = out_e3nn["energy"].sum()