diff --git a/src/MaxText/utils/ckpt_scripts/llama_mistral_mixtral_orbax_to_hf.py b/src/MaxText/utils/ckpt_scripts/llama_mistral_mixtral_orbax_to_hf.py
index 8a9afe34d..a1f525d29 100644
--- a/src/MaxText/utils/ckpt_scripts/llama_mistral_mixtral_orbax_to_hf.py
+++ b/src/MaxText/utils/ckpt_scripts/llama_mistral_mixtral_orbax_to_hf.py
@@ -48,7 +48,7 @@
 from transformers import LlamaForCausalLM, MistralForCausalLM, AutoModelForCausalLM, AutoConfig
 
 from MaxText import checkpointing
-from MaxText import llama_or_mistral_ckpt
+import llama_or_mistral_ckpt
 from MaxText import max_logging
 from MaxText import maxtext_utils
 from MaxText import pyconfig