Call hook to set OMP_NUM_THREADS

casparvl · Jul 25, 2024 · 00fca31 · 00fca31
1 parent d62443b
commit 00fca31
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/eessi/testsuite/tests/apps/PyTorch/PyTorch_torchvision.py b/eessi/testsuite/tests/apps/PyTorch/PyTorch_torchvision.py
@@ -69,6 +69,9 @@ def apply_setup_hooks(self):
         # This is a hybrid test, binding is important for performance
         hooks.set_compact_process_binding(self)
 
+        # Set OMP_NUM_THREADS based on the number of cores per task
+        hooks.set_omp_num_threads(self)
+
     @run_after('setup')
     def set_ddp_options(self):
         # Set environment variables for PyTorch DDP