From 1c15faee2adc7c7ec37933cc4fa9aac3a70813c8 Mon Sep 17 00:00:00 2001
From: michelle-yooh <yooh@google.com>
Date: Mon, 9 Oct 2023 22:08:41 +0000
Subject: [PATCH] add ici_tensor_parallelism=4 to test_decode

---
 end_to_end/test_decode.sh | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/end_to_end/test_decode.sh b/end_to_end/test_decode.sh
index c05d40ef6..0b32a38cf 100644
--- a/end_to_end/test_decode.sh
+++ b/end_to_end/test_decode.sh
@@ -17,6 +17,7 @@ fi
 #Train
 python3 MaxText/decode.py MaxText/configs/base.yml run_name=$RUN_NAME\
     steps=50 enable_checkpointing=False metrics_file='metrics.txt'\
-    base_output_directory=$OUTPUT_PATH dataset_path=$DATASET_PATH
+    base_output_directory=$OUTPUT_PATH dataset_path=$DATASET_PATH\
+    ici_tensor_parallelism=4
 
 python3 end_to_end/eval_assert.py metrics_average metrics.txt $NUM_TOKEN_THRESHOLD num_tokens