diff --git a/scripts/benchmark_sbatch_submit.sh b/scripts/benchmark_sbatch_submit.sh
index d0a67686cec8c3f677718f988dc466a8b599a424..6b5aa4e34d5ad7380ee9156cb8437fb5bcd92a82 100644
--- a/scripts/benchmark_sbatch_submit.sh
+++ b/scripts/benchmark_sbatch_submit.sh
@@ -1,41 +1,51 @@
-#!/bin/bash
+dim=2
+for nodes in {1..1}; do
+    for gpus in {1,8}; do
+        for batch_size in {64,128}; do
+            for iteration in {1..1}; do
 
-NUM_NODES=1
+                echo dim ${dim}, nodes ${nodes}, gpus ${gpus}, batch_size ${batch_size}, iteration ${iteration}
 
-dim=2
-if [ $NUM_NODES -eq 1 ]; then
-    for dim in {2,3}; do
-        for nodes in {1..1}; do
-            for gpus in {1,8}; do
-                for batch_size in {64,128}; do
-                    for iteration in {1..1}; do
+                # For single node
+                if [ $NUM_NODES -eq 1 ]; then
+                    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
+                else
 
-                        sleep 1 # pause to be kind to the scheduler
+                    # For multi node
+                    gpus=8
+                    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}
 
-                    done
-                done
+                    sleep 1 # pause to be kind to the scheduler
+                fi
             done
         done
     done
-else
-    for dim in {2,3}; do
-        for nodes in {2.. $NUM_NODES}; do
-            for gpus in {8,8}; do
-                for batch_size in {64,128}; do
-                    for iteration in {1..1}; do
+done
+
+dim=3
+for nodes in {1..1}; do
+    for gpus in {1,8}; do
+        for batch_size in {1,2,4}; do
+            for iteration in {1..1}; do
+
+                echo dim ${dim}, nodes ${nodes}, gpus ${gpus}, batch_size ${batch_size}, iteration ${iteration}
+
+                # For single node
+                if [ $NUM_NODES -eq 1 ]; then
+                    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}
 
-                        # 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
+                else
 
-                        sleep 1 # pause to be kind to the scheduler
+                    # For multi node
+                    gpus=8
+                    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}
 
-                    done
-                done
+                    sleep 1 # pause to be kind to the scheduler
+                fi
             done
         done
     done
-fi
+done