run_sweep: 2 threads for --device cpu jobs (matches Addendum 7 protocol)

1 files changed, 4 insertions, 1 deletions

diff --git a/scripts/run_sweep.py b/scripts/run_sweep.py
index 34cccf7..7b75921 100644
--- a/scripts/run_sweep.py
+++ b/scripts/run_sweep.py
@@ -60,7 +60,10 @@ def run_one(exp: str, job: dict, post_only: bool) -> dict:
     if not post_only:
         cmd = [PY, "-m", "src.train", job["model"], job["seed"],
                "--out_dir", f"runs/{exp}/{job['job']}"] + job["flags"]
-        env = dict(os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1")
+        # CPU-resident jobs get 2 threads (sequential tied cell benefits); GPU jobs
+        # stay at 1 so the feeder process doesn't contend with concurrent CPU jobs.
+        n_threads = "2" if "--device cpu" in job["flags"] else "1"
+        env = dict(os.environ, OMP_NUM_THREADS=n_threads, MKL_NUM_THREADS=n_threads)
         with open(logpath, "w") as lf:
             p = subprocess.run(cmd, cwd=ROOT, env=env, stdout=lf, stderr=subprocess.STDOUT)
         if p.returncode != 0: