also skip lammps dpa-2 tests for CUDA

Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Nov 12, 2024 · e8a99f4 · e8a99f4
1 parent 49f57bc
commit e8a99f4
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/source/lmp/tests/test_lammps_dpa_jax.py b/source/lmp/tests/test_lammps_dpa_jax.py
@@ -19,6 +19,12 @@
     write_lmp_data,
 )
 
+if os.environ.get("CUDA_VISIBLE_DEVICES", "") != "":
+    pytest.skip(
+        reason="The model is generated with CPU",
+        allow_module_level=True,
+    )
+
 pbtxt_file2 = (
     Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot-1.pbtxt"
 )