Enable the testing of CUDA Graph

openmm · Feb 25, 2022 · 6c22c88 · 6c22c88
1 parent 11ead20
commit 6c22c88
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/python/tests/TestTorchForce.py b/python/tests/TestTorchForce.py
@@ -8,7 +8,11 @@
 
 @pytest.mark.parametrize('platform', ['Reference', 'CUDA'])
 @pytest.mark.parametrize('precision', ['single', 'mixed', 'double'])
-def testEnergyForces(platform, precision):
+@pytest.mark.parametrize('useGraph', [True, False])
+def testEnergyForces(platform, precision, useGraph):
+
+    if pt.cuda.device_count() < 1 and platform == 'CUDA':
+        pytest.skip('A CUDA device is not available')
 
     # Create a system
     numParticles = 10
@@ -19,6 +23,8 @@ def testEnergyForces(platform, precision):
 
     # Create a TorchForce
     force = ot.TorchForce("../../tests/central.pt")
+    if useGraph:
+        force.setPlatformProperty('CUDAGraph', 'true')
     system.addForce(force)
 
     # Setup a simulation