diff --git a/tests/test_cueq.py b/tests/test_cueq.py index 21ec0e2e..5f47cbd9 100644 --- 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()