CHECK failed: (index) >= (0) when fine-tune MTF model about text-to-text-transfer-transformer HOT 6 CLOSED

t5_mesh_transformer  \
  --tpu="${TPU_NAME}" \
  --gcp_project="${PROJECT}" \
  --tpu_zone="${ZONE}" \
  --model_dir="${MODEL_DIR}/${EXP_NAME}" \
  --t5_tfds_data_dir="${DATA_DIR}" \
  --gin_file="dataset.gin" \
  --gin_param="utils.tpu_mesh_shape.model_parallelism = 1" \
  --gin_param="utils.tpu_mesh_shape.tpu_topology = '2x2'" \
  --use_model_api \
  --mode="finetune" \
  --train_steps="50000" \
  --mixture_or_task="${TASK_NAME}" \
  --pretrained_model_dir="${MODEL_DIR}/baseline-spm-fix-22000-unsupervised_mix_zh" \
  --checkpoint_mode="specific" \
  --checkpoint_steps="999600" \
  --gin_file="learning_rate_schedules/constant_0_001.gin" \
  --gin_param="MtfModel.save_checkpoints_steps = 5000" \
  --gin_param="MtfModel.batch_size = 128" \
  --gin_param="MtfModel.model_parallelism=1"

File "/usr/local/lib/python3.6/site-packages/mesh_tensorflow/transformer/dataset.py", line 513, in pack_dataset
    shapes = dataset.output_shapes
AttributeError: 'ShuffleDataset' object has no attribute 'output_shapes'
  In call to configurable 'pack_dataset' (<function pack_dataset at 0x7f317ac32ae8>)
  In call to configurable 'pack_or_pad' (<function pack_or_pad at 0x7f317ad39d90>)
  In call to configurable 'mesh_train_dataset_fn' (<function mesh_train_dataset_fn at 0x7f317ac442f0>)