MNT: cleanup in batch examples

elcorto · Nov 27, 2024 · 501d338 · 501d338
1 parent 2fa17c0
commit 501d338
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 8 deletions.
diff --git a/examples/batch_dask/run_psweep.py b/examples/batch_dask/run_psweep.py
@@ -29,16 +29,12 @@ def func(pset):
         local_directory="dask_tmp",
         log_directory="dask_log",
         scheduler_options={"dashboard_address": ":3333"},
-        # If you need a GPU
-        ##job_extra_directives=["--gres gpu:1"],
     )
 
     print(cluster.dashboard_link)
     print(cluster.job_script())
 
-    a = ps.plist("a", range(100))
-    ##params = ps.pgrid([a])
-    params = a
+    params = ps.plist("a", range(100))
 
     # Start 2 batch jobs, each with 10 dask workers (processes=10) and 10
     # cores, so 1 core (1 thread) / worker and 20 workers in total (2 jobs x 10

diff --git a/examples/batch_dask/run_psweep_jax_gpu.py b/examples/batch_dask/run_psweep_jax_gpu.py
@@ -189,9 +189,7 @@ def func(pset):
         scheduler_options={"dashboard_address": ":3333"},
     )
 
-    nn = ps.plist("nn", ps.intspace(100, 300, num=100))
-    ##params = ps.pgrid([nn])
-    params = nn
+    params = ps.plist("nn", ps.intspace(100, 300, num=100))
 
     cluster.scale(jobs=3)
     client = Client(cluster)