Add _mp_fn to run_mae.py for XLA testing
Testing MAE ViT on TPU https://github.com/GoogleCloudPlatform/ml-testing-accelerators/pull/771

Based on update required to `run_glue.py` in https://github.com/huggingface/transformers/pull/4146
_The documentation is not available anymore as the PR was closed or merged._