set CUDA_VISIBLE_DEVICES based on workerID

Libensemble · Jul 12, 2021 · 7dee9a1 · 7dee9a1
1 parent c19749c
commit 7dee9a1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/libensemble/tests/scaling_tests/ddmd/keras_cvae_ml_genf.py b/libensemble/tests/scaling_tests/ddmd/keras_cvae_ml_genf.py
@@ -174,6 +174,7 @@ def run_keras_cvae_ml_genf(H, persis_info, gen_specs, libE_info):
     exctr = Executor.executor
     persis_info['stage_count'] = -1
     os.environ["OMP_NUM_THREADS"] = '4'
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(libE_info['workerID'])
     initial_complete = False
     tag = None
 

diff --git a/libensemble/tests/scaling_tests/ddmd/openmm_md_simf.py b/libensemble/tests/scaling_tests/ddmd/openmm_md_simf.py
@@ -49,6 +49,7 @@ def run_openmm_sim_f(H, persis_info, sim_specs, libE_info):
     """
     calc_status = 0
     os.environ["OMP_NUM_THREADS"] = '4'
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(libE_info['workerID'])
 
     #  Specify the Executor object created in the calling script.
     exctr = Executor.executor