Merge pull request #1442 from AI-Hypercomputer:wang2yn84-patch-4

PiperOrigin-RevId: 743763850

Author: maxtext authors

MaxText/train.py

@@ -948,7 +948,6 @@ def train_loop(config, state=None):
 def main(argv: Sequence[str]) -> None:
   pathwaysutils.initialize()
   jax.config.update("jax_default_prng_impl", "unsafe_rbg")
-  jax.config.update("jax_enable_compilation_cache", os.environ.get("JAX_ENABLE_COMPILATION_CACHE", True))
   # TF allocates extraneous GPU memory when using TFDS data
   # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF
   tf.config.set_visible_devices([], "GPU")