Merge pull request #2680 from AI-Hypercomputer:mohit/fix_ckpt

PiperOrigin-RevId: 832027535

Author: Google-ML-Automation

src/MaxText/utils/ckpt_conversion/to_maxtext.py

@@ -61,7 +61,7 @@
 import threading
 from functools import partial
 from typing import Sequence, List, Any, Callable
-from benchmarks.benchmark_utils import str2bool
+from MaxText.inference_utils import str2bool
 import numpy as np
 import jax
 import psutil