From 2a24b906ac9139b57ee6c0282f76c148b70b594d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Wed, 27 Jul 2022 14:36:22 +0200
Subject: [PATCH] Add batch size script argument for standalone tests (#13841)

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
---
 tests/tests_pytorch/run_standalone_tests.sh | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/tests/tests_pytorch/run_standalone_tests.sh b/tests/tests_pytorch/run_standalone_tests.sh
index 45264bb74a..5297cbd033 100644
--- a/tests/tests_pytorch/run_standalone_tests.sh
+++ b/tests/tests_pytorch/run_standalone_tests.sh
@@ -15,6 +15,20 @@
 set -e
 # THIS FILE ASSUMES IT IS RUN INSIDE THE tests/tests_pytorch DIRECTORY
 
+# Batch size for testing: Determines how many standalone test invocations run in parallel
+test_batch_size=6
+
+while getopts "b:" opt; do
+    case $opt in
+        b)
+            test_batch_size=$OPTARG;;
+        *)
+            echo "Usage: $(basename $0) [-b batch_size]"
+            exit 1;;
+    esac
+done
+shift $((OPTIND-1))
+
 # this environment variable allows special tests to run
 export PL_RUN_STANDALONE_TESTS=1
 # python arguments
@@ -40,7 +54,6 @@ parametrizations_arr=($parametrizations)
 # tests to skip - space separated
 blocklist='profilers/test_profiler.py::test_pytorch_profiler_nested_emit_nvtx utilities/test_warnings.py'
 report=''
-test_batch_size=6
 
 rm -f standalone_test_output.txt  # in case it exists, remove it
 function show_batched_output {