Error when running the T5 pre-training script
Author: hugoabonizio
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.15.0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): not installed (NA)
- Tensorflow version (GPU?): 2.7.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.6 (tpu)
- Jax version: 0.2.26
- JaxLib version: 0.1.75
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help
@patrickvonplaten @patil-suraj @sgugger

## Information

I'm running the `run_t5_mlm_flax.py` example script on a custom version of mC4 dataset (only English), but I'm facing this exception during the training after a few steps (changing the seed changes the step that it occours, as expected).

```
Traceback (most recent call last):
  File "/home/hugo/run_t5_mlm_flax.py", line 925, in <module>
    main()
  File "/home/hugo/run_t5_mlm_flax.py", line 839, in main
    model_inputs = data_collator(samples)
  File "/home/hugo/run_t5_mlm_flax.py", line 339, in __call__
    batch["input_ids"] = self.filter_input_ids(input_ids, input_ids_sentinel)
  File "/home/hugo/run_t5_mlm_flax.py", line 392, in filter_input_ids
    input_ids = input_ids_full[input_ids_full > 0].reshape((batch_size, -1))
ValueError: cannot reshape array of size 65471 into shape (64,newaxis)
```

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

I'm running the script with the following parameters:

```
python $HOME/run_t5_mlm_flax.py \
	--output_dir="./models/${MODEL_DIR}" \
	--overwrite_output_dir \
	--model_type="t5" \
	--config_name="./models/${MODEL_DIR}" \
	--tokenizer_name="google/byt5-small" \
	--train_file="/mnt/datasets/train/en.json" \
	--max_seq_length="1024" \
	--num_train_epochs="1" \
	--per_device_train_batch_size="8" \
	--per_device_eval_batch_size="8" \
	--adafactor \
	--learning_rate="0.001" \
	--weight_decay="0.001" \
	--warmup_steps="2000" \
	--save_steps="50000" \
	--eval_steps="50000" \
	--push_to_hub \
	--logging_steps="1000" \
	--preprocessing_num_workers="96" \
	--gradient_accumulation_steps="128"
```

## Expected behavior

Not throwing this exception.

Comments:
Hey @hugoabonizio , could you please have a look at:

https://github.com/huggingface/transformers/issues/12755

I've seen the same issue, my current workaround is to try-catch `data_collator` calls:

```diff
@@ -813,13 +818,17 @@ def main():
 
         # Gather the indexes for creating the batch and do a training step
         for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1)):
-            samples = [tokenized_datasets["train"][int(idx)] for idx in batch_idx]
-            model_inputs = data_collator(samples)
-
-            # Model forward
-            model_inputs = shard(model_inputs.data)
-            state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
-            train_metrics.append(train_metric)
+            try:
+                samples = [tokenized_datasets["train"][int(idx)] for idx in batch_idx]
+                model_inputs = data_collator(samples)
+
+                # Model forward
+                model_inputs = shard(model_inputs.data)
+                state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
+                train_metrics.append(train_metric)
+            except:
+                cur_step = epoch * (num_train_samples // train_batch_size) + step
+                continue
```

Hi @stefan-it, sorry, this issue sliped under my radar :sweat_smile: 

Thank you for the code snippet, I was thinking on a similar workaround but will use yours.


Sorry haven't found the time to take a closer look. Will try to do it tomorrow!

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

@hugoabonizio  - could you check whether this PR solves the problem: https://github.com/huggingface/transformers/pull/15835 ? :-) Sorry for being so late here

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

@patrickvonplaten sorry for the late reply. It looks like it's fixed now, thank you!