add ici_tensor_parallelism=4 to test_decode

AI-Hypercomputer · Oct 9, 2023 · 1c15fae · 1c15fae
1 parent 1ca17b3
commit 1c15fae
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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