Flax MLM example script has PyTorch dependency
**Script in question:** [run_mlm_flax.py](https://github.com/huggingface/transformers/blob/master/examples/flax/language-modeling/run_mlm_flax.py)

```bash
Traceback (most recent call last):
  File "./run_mlm_flax.py", line 319, in <module>
    f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
  File "/home/christopher/transformers/src/transformers/file_utils.py", line 1641, in wrapper
    raise ImportError(f"Method `{func.__name__}` requires PyTorch.")
ImportError: Method `device` requires PyTorch.
```

## 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.9.0.dev0' (editable pip install from `git` source)
- Platform: TPU v3-8 (Operating System: Ubuntu 20.04.2 LTS, Kernel: Linux 5.4.0-1043-gcp, Architecture: x86-64)
- Python version: Python 3.8.10 [GCC 9.4.0] on linux
- PyTorch version (GPU?): '1.9.0+cu102' (extras['all'])
- Tensorflow version (GPU?): '2.5.0' (extras['all'])
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: Yes

### Who can help

@patrickvonplaten and @patil-suraj 

## Information

Model I am using: roBERTa

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:

1. Set up a Flax-based Huggingface environment (`$ pip install -e ".[flax]"`) according to step 4 [here](https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects#tpu-vm)
2. Follow [the tutorial](https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects#how-to-use-the-hub-for-collaboration) on training roBERTa on Alemannic OSCAR.


## Expected behavior
 The script should run with no error. The absence of PyTorch should have no bearing on a Flax script.