From e88da56131593f2c5ea31d49936261bce4d562ee Mon Sep 17 00:00:00 2001
From: Karan Goel <karangoel@google.com>
Date: Thu, 24 Apr 2025 14:15:36 -0700
Subject: [PATCH] model_ckpt_conversion: use python module, enable gcloud
 checkpoint downloading

---
 jetstream/tools/maxtext/model_ckpt_conversion.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/jetstream/tools/maxtext/model_ckpt_conversion.sh b/jetstream/tools/maxtext/model_ckpt_conversion.sh
index 7c135175..7017fb69 100644
--- a/jetstream/tools/maxtext/model_ckpt_conversion.sh
+++ b/jetstream/tools/maxtext/model_ckpt_conversion.sh
@@ -65,11 +65,11 @@ else
     pip install torch --index-url https://download.pytorch.org/whl/cpu
     # llama_or_mistral_ckpt.py requires local path, so we need to copy the checkpoint from CHKPT_BUCKET to local.
     tmp_ckpt_path="/tmp/"
-    #gcloud storage cp -r ${CHKPT_BUCKET} ${tmp_ckpt_path}
+    gcloud storage cp -r ${CHKPT_BUCKET} ${tmp_ckpt_path}
 
     path_parts=(${CHKPT_BUCKET//\// })
     directory_substring=${path_parts[-1]}
-    CONVERT_CKPT_SCRIPT="llama_or_mistral_ckpt.py"
+    CONVERT_CKPT_SCRIPT="llama_or_mistral_ckpt"
 
     if [[ ! -z "${LORA_INPUT_ADAPTERS_PATH}" ]]; then
 	lora_local_path="/tmp/"