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