try fix

tcclevenger · Nov 18, 2023 · f0af467 · f0af467
1 parent c60716d
commit f0af467
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/core/unit_test/TestDeviceAndThreads.py b/core/unit_test/TestDeviceAndThreads.py
@@ -17,6 +17,7 @@
 
 import unittest
 import subprocess
+import psutil
 
 PREFIX = "$<TARGET_FILE_DIR:Kokkos_CoreUnitTest_DeviceAndThreads>"
 EXECUTABLE = "$<TARGET_FILE_NAME:Kokkos_CoreUnitTest_DeviceAndThreads>"
@@ -30,7 +31,9 @@ def GetFlag(flag, *extra_args):
     return int(p.stdout)
 
 def GetNumThreads(max_threads):
-    for x in [1, 2, 4, 6, 8]:
+    phys_cores_count = psutil.cpu_count(logical=False)
+    looplist = [1] + [i*phys_cores_count for i in [1,2,3,4,5,6,7]]
+    for x in looplist:
         if x >= max_threads:
             break
         yield x