diff --git a/README.md b/README.md
index 2c3ec9558a32dfe8920769077c5af09b4bb8c3fc..afc9592a581c1f80c74aa35d0d16241e42d0472d 100644
--- a/README.md
+++ b/README.md
@@ -59,14 +59,16 @@ bash scripts/copy_data_for_benchmark.sh
 You can choose either singularity or enroot in the scripts ```benchmark_single_node.sbatch``` and ```benchmark_multi_node.sbatch```.
 
 The input arguments are: 
-1. Number of nodes,
-2. Number of gpus per node,
-3. Number of iterations for each parameter setting.
+1. Data dimention,
+2. Number of nodes,
+3. Number of gpus per node,
+4. Number of iterations for each parameter setting.
+5. Batch size
 
 We will average the benchmark performance over the iterations. The maximum usable (without a OOM error) batch size is 256 and 128 for single and multi-node, respectively.
 ```
 cd cd Berzelius-nnU-Net-Benchmark && mkdir -p sbatch_out
-bash scripts/benchmark_sbatch_submit.sh 1 8 100 128
+bash scripts/benchmark_sbatch_submit.sh 2 1 8 10 128
 ```
 
 ### Results  
diff --git a/scripts/benchmark_sbatch_submit.sh b/scripts/benchmark_sbatch_submit.sh
index da5b934c3f2d3e3969b77ffca7eab0b5583f3097..f077dd0066bc761b4fa9127d9558f1148f8818e1 100644
--- a/scripts/benchmark_sbatch_submit.sh
+++ b/scripts/benchmark_sbatch_submit.sh
@@ -1,40 +1,35 @@
 #!/bin/bash
 
-mkdir -p sbatch_out
+NUM_NODES=$2
+NUM_GPUS=$3
+NUM_ITERATIONS=$4
 
-dim=2
-NUM_NODES=$1
-NUM_GPUS=$2
-NUM_ITERATIONS=$3
-BATCH_SIZE=$4
+dim=$1
+batch_size=$5
 if [ $NUM_NODES -eq 1 ]; then
-    for nodes in {1..$((NUM_NODES))}; do
-        for gpus in {1..$((NUM_GPUS))}; do
-            for batch_size in $((BATCH_SIZE)); do
-                for iteration in {1..$((NUM_ITERATIONS))}; do
+    for ((nodes = 1; nodes <= $NUM_NODES; nodes++)); do
+        for ((gpus = 1; gpus <= $NUM_GPUS; gpus++)); do
+            for ((iteration = 1; iteration <= $NUM_ITERATIONS; iteration++)); do
+                
+                echo dim ${dim}, nodes ${nodes}, gpus ${gpus}, batch_size ${batch_size}, iteration ${iteration}
+                # For single node
+                #sbatch -o sbatch_out/benchmark_dim${dim}_nodes${nodes}_gpus${gpus}_batchsize${batch_size}_iteration${iteration}.out scripts/benchmark_single_node.sbatch ${dim} ${nodes} ${gpus} ${batch_size} ${iteration}
 
-                    echo dim ${dim}, nodes ${nodes}, gpus ${gpus}, batch_size ${batch_size}, iteration ${iteration}
-                    # For single node
-                    #sbatch -o sbatch_out/benchmark_dim${dim}_nodes${nodes}_gpus${gpus}_batchsize${batch_size}_iteration${iteration}.out scripts/benchmark_single_node.sbatch ${dim} ${nodes} ${gpus} ${batch_size} ${iteration}
+                sleep 1 # pause to be kind to the scheduler
 
-                    sleep 1 # pause to be kind to the scheduler
-
-                done
             done
         done
     done
 else
-    for nodes in {2..$((NUM_NODES))}; do
-        for gpus in {$((NUM_GPUS))}; do
-            for batch_size in $((BATCH_SIZE)); do
-                for iteration in {1..$((NUM_ITERATIONS))}; do
-
-                    # For multi node
-                    sbatch -o sbatch_out/benchmark_dim${dim}_nodes${nodes}_gpus${gpus}_batchsize${batch_size}_iteration${iteration}.out scripts/benchmark_multi_node.sbatch ${dim} ${nodes} ${gpus} ${batch_size} ${iteration}
+    for ((nodes = 2; nodes <= $NUM_NODES; nodes++)); do
+        for ((gpus = 8; gpus <= $NUM_GPUS; gpus++)); do
+            for ((iteration = 1; iteration <= $NUM_ITERATIONS; iteration++)); do
+                batch_size=$BATCH_SIZE
+                # For multi node
+                sbatch -o sbatch_out/benchmark_dim${dim}_nodes${nodes}_gpus${gpus}_batchsize${batch_size}_iteration${iteration}.out scripts/benchmark_multi_node.sbatch ${dim} ${nodes} ${gpus} ${batch_size} ${iteration}
 
-                    sleep 1 # pause to be kind to the scheduler
-                done
+                sleep 1 # pause to be kind to the scheduler
             done
         done
     done
-fi
+fi
\ No newline at end of file