fix test_relaxation

tests/test_relaxation.py

@@ -32,7 +32,11 @@ def test_relaxation(
     chgnet.graph_converter = converter
     relaxer = StructOptimizer(model=chgnet)
     result = relaxer.relax(
-        li_mn_o2, verbose=True, ase_filter=ase_filter, assign_magmoms=assign_magmoms
+        li_mn_o2,
+        verbose=True,
+        ase_filter=ase_filter,
+        assign_magmoms=assign_magmoms,
+        steps=10,  # few steps for faster testing
     )
     assert list(result) == ["final_structure", "trajectory"]
     final_struct, traj = result["final_structure"], result["trajectory"]
@@ -50,10 +54,10 @@ def test_relaxation(
     assert {*traj.__dict__} == {
         *"atoms energies forces stresses magmoms atom_positions cells".split()
     }
-    assert len(traj) == 2 if algorithm == "legacy" else 4
+    assert len(traj) == 12, f"{len(traj)=}, {algorithm=}"
 
     # make sure final structure is more relaxed than initial one
-    assert traj.energies[-1] == pytest.approx(-58.94209, rel=1e-4)
+    assert traj.energies[-1] == pytest.approx(-58.6675376, abs=0.3)
 
 
 no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device")