XLNET on SQuAD2 evaluation error
## 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.7.0.dev0
- Platform: NAME="Red Hat Enterprise Linux Server" VERSION="7.9 (Maipo)"
- Python version: Python 3.6.8 :: Anaconda, Inc.
- PyTorch version (GPU?):  1.8.1
- Tensorflow version (GPU?): not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: parallel in 4 GPUs on one node

### Who can help

@patrickvonplaten @sgugger 
- Models: XLNET, benchmarks on SQuAD2

## Information

Model I am using (Bert, XLNet ...): XLNET cased-base, it produces error like following in the evaluation step:
```
100%|█████████▉| 1528/1529 [09:59<00:00,  2.55it/s]Traceback (most recent call last):
  File "/scratch365/yding4/Amazon_2021_summer_intern/AVEQA_PyTorch/XLNET_on_SQuAD2.0/run_qa_beam_search.py", line 661, in <module>
    main()
  File "/scratch365/yding4/Amazon_2021_summer_intern/AVEQA_PyTorch/XLNET_on_SQuAD2.0/run_qa_beam_search.py", line 620, in main
    metrics = trainer.evaluate()
  File "/scratch365/yding4/Amazon_2021_summer_intern/AVEQA_PyTorch/XLNET_on_SQuAD2.0/trainer_qa.py", line 50, in evaluate
    ignore_keys=ignore_keys,
  File "/scratch365/yding4/Amazon_2021_summer_intern/transformers/src/transformers/trainer.py", line 2169, in evaluation_loop
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "/scratch365/yding4/Amazon_2021_summer_intern/transformers/src/transformers/trainer.py", line 2383, in prediction_step
    outputs = model(**inputs)
  File "/afs/crc.nd.edu/user/y/yding4/.conda/envs/AVEQA_PyTorch/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/afs/crc.nd.edu/user/y/yding4/.conda/envs/AVEQA_PyTorch/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py", line 168, in forward
    return self.gather(outputs, self.output_device)
  File "/afs/crc.nd.edu/user/y/yding4/.conda/envs/AVEQA_PyTorch/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py", line 180, in gather
    return gather(outputs, output_device, dim=self.dim)
  File "/afs/crc.nd.edu/user/y/yding4/.conda/envs/AVEQA_PyTorch/lib/python3.6/site-packages/torch/nn/parallel/scatter_gather.py", line 76, in gather
    res = gather_map(outputs)
  File "/afs/crc.nd.edu/user/y/yding4/.conda/envs/AVEQA_PyTorch/lib/python3.6/site-packages/torch/nn/parallel/scatter_gather.py", line 70, in gather_map
    for k in out))
  File "<string>", line 11, in __init__
  File "/scratch365/yding4/Amazon_2021_summer_intern/transformers/src/transformers/file_utils.py", line 1739, in __post_init__
    for element in iterator:
  File "/afs/crc.nd.edu/user/y/yding4/.conda/envs/AVEQA_PyTorch/lib/python3.6/site-packages/torch/nn/parallel/scatter_gather.py", line 70, in <genexpr>
    for k in out))
  File "/afs/crc.nd.edu/user/y/yding4/.conda/envs/AVEQA_PyTorch/lib/python3.6/site-packages/torch/nn/parallel/scatter_gather.py", line 71, in gather_map
    return type(out)(map(gather_map, zip(*outputs)))
  File "/afs/crc.nd.edu/user/y/yding4/.conda/envs/AVEQA_PyTorch/lib/python3.6/site-packages/torch/nn/parallel/scatter_gather.py", line 63, in gather_map
    return Gather.apply(target_device, dim, *outputs)
  File "/afs/crc.nd.edu/user/y/yding4/.conda/envs/AVEQA_PyTorch/lib/python3.6/site-packages/torch/nn/parallel/_functions.py", line 72, in forward
    return comm.gather(inputs, ctx.dim, ctx.target_device)
  File "/afs/crc.nd.edu/user/y/yding4/.conda/envs/AVEQA_PyTorch/lib/python3.6/site-packages/torch/nn/parallel/comm.py", line 235, in gather
    return torch._C._gather(tensors, dim, destination)
RuntimeError: Input tensor at index 3 has invalid shape [384, 1, 1024], but expected [384, 2, 1024]
```

The problem arises when using:
- [ * ] the official example scripts: (give details below)

```
python run_qa_beam_search.py \
    --model_name_or_path xlnet-large-cased \
    --dataset_name squad_v2 \
    --do_train \
    --do_eval \
    --version_2_with_negative \
    --learning_rate 3e-5 \
    --num_train_epochs 1 \
    --max_seq_length 384 \
    --doc_stride 128 \
    --output_dir ./wwm_cased_finetuned_squad/ \
    --per_device_eval_batch_size=2  \
    --per_device_train_batch_size=2   \
    --save_steps 5000
```
The evaluation works when the **--per_device_eval_batch_size=2** is set to 1.
The tasks I am working on is:
- [x] an official GLUE/SQUaD task: SQUaD2.0 task
the script is used from https://github.com/huggingface/transformers/blob/master/examples/pytorch/question-answering/run_qa_beam_search.py

## To reproduce

Steps to reproduce the behavior:

1. just run the script


## Expected behavior
train successfully and get evaluation results
It looks like you are using DataParallel for the evaluation, which does not always work if the number of samples is not a round multiple of the number of GPUs. You should use DistributedDataParallel (as recommended by the PyTorch team), just change your command to:
```
python -m torch.distributed.launch --nproc_per_node 2 run_qa_beam_search.py \
``` 
(and replace 2 by your actual number of GPUs) and it should work (it did on my setup).<|||||>> It looks like you are using DataParallel for the evaluation, which does not always work if the number of samples is not a round multiple of the number of GPUs. You should use DistributedDataParallel (as recommended by the PyTorch team), just change your command to:
> 
> ```
> python -m torch.distributed.launch --nproc_per_node 2 run_qa_beam_search.py \
> ```
> 
> (and replace 2 by your actual number of GPUs) and it should work (it did on my setup).

Thank you so much, close the issue now.