new Transformer update causes an error with TPU XLA implementation
Hi,
I notice the new release of the Transformer model (4.20) causes an issue with PyTorch XLA implementation and the new error message says "Cannot replicate if number of devices (1) is different from 8" appears.



### Who can help?

_No response_

### Information

- [x] The official example scripts
- [ ] 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

`!pip install cloud-tpu-client==0.10 torch==1.11.0 https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-1.11-cp37-cp37m-linux_x86_64.whl`

```
!pip3 install git+https://github.com/huggingface/transformers
!git clone https://github.com/huggingface/transformers
```


`!pip3 install -r /content/transformers/examples/pytorch/question-answering/requirements.txt`

```
!python /content/transformers/examples/pytorch/xla_spawn.py --num_cores=8 /content/transformers/examples/pytorch/question-answering/run_qa.py --model_name_or_path sultan/BioM-ELECTRA-Large-Discriminator \
--dataset_name squad_v2 \
--do_train \
--do_eval \
--dataloader_num_workers 4 \
--preprocessing_num_workers 4 \
--version_2_with_negative \
--num_train_epochs 2 \
--learning_rate 5e-5 \
--max_seq_length 384 \
--doc_stride 128 \
--per_device_train_batch_size 4 \
--per_device_eval_batch_size 4 \
--logging_steps 500 \
--save_steps 1000 \
--overwrite_output_dir \
--output_dir out
```


To fix the issue we git clone and git+install with 4.19 release:

```
!pip3 install git+https://github.com/huggingface/transformers.git@v4.19.4
!git clone --depth 1 --branch v4.19.4 https://github.com/huggingface/transformers
```


### Expected behavior

```shell
A new error message says "Cannot replicate if number of devices (1) is different from 8" appears.
```

Maybe of interest to @sgugger <|||||>cc @muellerzr It's probably linked to your recent changes for selecting the TPU device.