diff --git a/README.md b/README.md
index 2314b50aa74e1a1c7dd105db52b6fd3b221c6968..8f9488c05795a63fd22114988fc082e278eebd9d 100644
--- a/README.md
+++ b/README.md
@@ -83,9 +83,14 @@ 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 argument is the number of nodes used for benchmarking.
+The input arguments are: 
+1. Number of nodes,
+2. Number of gpus per node,
+3. Number of iterations for each parameter setting.
+
+We will average the benchmark performance over the iterations. We use a batch size of 128 which is the maximum usable batch size without a OOM error.
 ```
-bash scripts/benchmark_sbatch_submit.sh 1
+bash scripts/benchmark_sbatch_submit.sh 1 8 100
 ```
 
 ### Results  
diff --git a/scripts/benchmark_sbatch_submit.sh b/scripts/benchmark_sbatch_submit.sh
index 75e7d2dc0af03e457ca5a6969c5076f1a3ce0b97..44ac79bc2fde0e8acea54650fdd764b79e621658 100644
--- a/scripts/benchmark_sbatch_submit.sh
+++ b/scripts/benchmark_sbatch_submit.sh
@@ -1,12 +1,16 @@
+#!/bin/bash
+
 mkdir -p sbatch_out
 
 dim=2
 NUM_NODES=$1
+NUM_GPUS=$2
+NUM_ITERATIONS=$3
 if [ $NUM_NODES -eq 1 ]; then
     for nodes in {1..1}; do
-        for gpus in {1..8}; do
+        for gpus in {1..$NUM_GPUS}; do
             for batch_size in 128; do
-                for iteration in {1..100}; do
+                for iteration in {1..$NUM_ITERATIONS}; do
 
                     echo dim ${dim}, nodes ${nodes}, gpus ${gpus}, batch_size ${batch_size}, iteration ${iteration}
                     # For single node
@@ -20,9 +24,9 @@ if [ $NUM_NODES -eq 1 ]; then
     done
 else
     for nodes in {2..$NUM_NODES}; do
-        for gpus in {8}; do
+        for gpus in {$NUM_GPUS}; do
             for batch_size in 128; do
-                for iteration in {1..100}; 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}