fix: increase batch

based-robotics · Sep 17, 2024 · a24ad63 · a24ad63
1 parent 3c5bd18
commit a24ad63
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/benchmarks/run_benchmark.py b/benchmarks/run_benchmark.py
@@ -43,11 +43,11 @@ class ColabPaths:
 
     torch.manual_seed(0)
 
-N_ENVS_SWEEP = [1, 2, 4, 8, 16, 32, 64, 128, 256]  # , 512, 1024, 2048, 4096, 8192, 16384, 32768]
+N_ENVS_SWEEP = [2**i for i in range(20)]
 N_EVALS = 20
 
 # Load functions for CUDA benchmarking
-fn_files = ["fn_1e1.casadi", "fn_1e2.casadi", "fn_1e3.casadi"]  # , "fn_1e4.casadi", "fn_1e5.casadi"]
+fn_files = ["fn_1e1.casadi", "fn_1e2.casadi"]
 benchmark_fns = [Function.load(os.path.join(PathsProvider.FUNCTIONS_DIR, fn)) for fn in fn_files]