Run Vit-MAE script key error
### System Info

- `transformers` version: 4.23.1
- Platform: Linux-5.10.133+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.15
- Huggingface_hub version: 0.10.1
- PyTorch version (GPU?): 1.12.1+cu113 (True)
- Tensorflow version (GPU?): 2.9.2 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@NielsRogge @sg

### Information

- [x] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

I've put down the main parts of the run MAE script [https://github.com/huggingface/transformers/blob/main/examples/pytorch/image-pretraining/run_mae.py]
into a colab notebook: 
[https://colab.research.google.com/drive/1WtOTp-ocbBTgVXFiEXuY8MWhz_ex9OBy?usp=sharing]
When I reach the trainer.train() step I get a key error from the data collator function - is the script perhaps outdated, or am I doing something wrong?

### Expected behavior

Trainer starts training
You'll need to add `remove_unused_columns=False` to the `TrainingArguments`.

The reason is because of the use of `set_transform` when preparing the datasets, which does things on-the-fly. Hence we still need to `image` column in the datasets to turn them into `pixel_values`.<|||||>Amazing thank you!