diff --git a/docs/usage/eval.md b/docs/usage/eval.md
index e29a42754..3b8ab2f7f 100644
--- a/docs/usage/eval.md
+++ b/docs/usage/eval.md
@@ -168,7 +168,7 @@ cause issues), run the following on commandline:
 
 ```sh
 EVAL_OUTPUT_DIR="/tmp/model-eval/"
-python -m t5x.eval \
+python -m t5x.eval_unfragmented \
   --gin_file=t5x/google/examples/flaxformer_t5/configs/examples/eval/t5_1_1_small_cbqa_natural_questions.gin \
   --gin.EVAL_OUTPUT_DIR=\"${EVAL_OUTPUT_DIR}\" \
   --alsologtostderr
diff --git a/docs/usage/finetune.md b/docs/usage/finetune.md
index 5969e8f0e..5a27f16f3 100644
--- a/docs/usage/finetune.md
+++ b/docs/usage/finetune.md
@@ -161,7 +161,7 @@ cause issues), run the following on commandline:
 
 ```sh
 MODEL_DIR="/tmp/finetune-model/"
-python -m t5x.train \
+python -m t5x.train_unfragmented \
   --gin_file=t5x/examples/t5/t5_1_1/examples/small_wmt_finetune.gin \
   --gin.MODEL_DIR=\"${MODEL_DIR}\" \
   --alsologtostderr
diff --git a/docs/usage/gin.md b/docs/usage/gin.md
index 0c6bce3f6..035a603bb 100644
--- a/docs/usage/gin.md
+++ b/docs/usage/gin.md
@@ -335,7 +335,7 @@ and if you are doing finetuning, you can use
 We can apply these two files with the following command:
 
 ```sh
-python -m t5x.train \
+python -m t5x.train_unfragmented \
   --gin_file=t5x/examples/t5/t5_1_1/small.gin \
   --gin_file=t5x/configs/runs/finetune.gin \
   --logtostderr
@@ -354,7 +354,7 @@ you'll see the list of required overrides, which we will populate for finetuning
 on WMT in the updated launch command here:
 
 ```sh
-python -m t5x.train \
+python -m t5x.train_unfragmented \
   --gin_file=t5x/examples/t5/t5_1_1/small.gin \
   --gin_file=t5x/configs/runs/finetune.gin \
   --gin.MIXTURE_OR_TASK_NAME=\"wmt_t2t_ende_v003\" \
@@ -379,7 +379,7 @@ is just such an example that allows you to exactly duplicate the previous launch
 command simply by calling:
 
 ```sh
-python -m t5x.train \
+python -m t5x.train_unfragmented \
   --gin_file=t5x/examples/t5/t5_1_1/examples/small_wmt_finetune.gin \
   --gin.MODEL_DIR=\"/tmp/t5_1_1_small_finetune_gin\" \
   --logtostderr
diff --git a/docs/usage/infer-files.md b/docs/usage/infer-files.md
index a08165b6b..95d88c25c 100644
--- a/docs/usage/infer-files.md
+++ b/docs/usage/infer-files.md
@@ -161,7 +161,7 @@ cause issues), run the following on commandline:
 
 ```sh
 INFER_OUTPUT_DIR="/tmp/model-infer/"
-python -m t5x.infer \
+python -m t5x.infer_unfragmented \
   --gin_file=t5x/google/examples/flaxformer_t5/configs/examples/inference/t5_1_1_small_cbqa_natural_questions_tfexample.gin \
   --gin.INFER_OUTPUT_DIR=\"${INFER_OUTPUT_DIR}\" \
   --alsologtostderr
diff --git a/docs/usage/infer-seqio.md b/docs/usage/infer-seqio.md
index ecc5cdec9..541d52dee 100644
--- a/docs/usage/infer-seqio.md
+++ b/docs/usage/infer-seqio.md
@@ -159,7 +159,7 @@ cause issues), run the following on commandline:
 
 ```sh
 INFER_OUTPUT_DIR="/tmp/model-infer/"
-python -m t5x.infer \
+python -m t5x.infer_unfragmented \
   --gin_file=t5x/google/examples/flaxformer_t5/configs/examples/inference/t5_1_1_small_cbqa_natural_questions.gin \
   --gin.INFER_OUTPUT_DIR=\"${INFER_OUTPUT_DIR}\" \
   --alsologtostderr
diff --git a/docs/usage/pretrain.md b/docs/usage/pretrain.md
index c302974b6..219647111 100644
--- a/docs/usage/pretrain.md
+++ b/docs/usage/pretrain.md
@@ -140,7 +140,7 @@ cause issues), run the following on commandline:
 
 ```sh
 MODEL_DIR="/tmp/pretrain-model/"
-python -m t5x.train \
+python -m t5x.train_unfragmented \
   --gin_file=t5x/examples/t5/t5_1_1/c4_pretrain_small.gin \
   --gin.MODEL_DIR=\"${MODEL_DIR}\" \
   --alsologtostderr
diff --git a/t5x/notebooks/evaluation.ipynb b/t5x/notebooks/evaluation.ipynb
index 490cb2451..ebb9add07 100644
--- a/t5x/notebooks/evaluation.ipynb
+++ b/t5x/notebooks/evaluation.ipynb
@@ -870,7 +870,7 @@
         "\n",
         "```\n",
         "EVAL_OUTPUT_DIR=\"/tmp/eval-model/\"\n",
-        "python -m t5x.train \\\n",
+        "python -m t5x.train_unfragmented \\\n",
         "  --gin_file=${GIN_FILE_PATH} \\\n",
         "  --gin.EVAL_OUTPUT_DIR=\\\"${EVAL_OUTPUT_DIR}\\\" \\\n",
         "  --alsologtostderr\n",
diff --git a/t5x/notebooks/inference.ipynb b/t5x/notebooks/inference.ipynb
index d1b84b710..7706ab958 100644
--- a/t5x/notebooks/inference.ipynb
+++ b/t5x/notebooks/inference.ipynb
@@ -645,7 +645,7 @@
         "\n",
         "```\n",
         "INFER_OUTPUT_DIR=\"/tmp/inference-model/\"\n",
-        "python -m t5x.infer \\\n",
+        "python -m t5x.infer_unfragmented \\\n",
         "  --gin_file=${GIN_FILE_PATH} \\\n",
         "  --gin.INFER_OUTPUT_DIR=\\\"${INFER_OUTPUT_DIR}\\\" \\\n",
         "  --alsologtostderr\n",
diff --git a/t5x/train.py b/t5x/train.py
index 0ab63bdaf..06cd0a32b 100644
--- a/t5x/train.py
+++ b/t5x/train.py
@@ -16,6 +16,7 @@
 
 """
 
+
 # pylint: disable=g-import-not-at-top
 
 import functools