`trainer.predict(dataset)` drops samples when used with multi-gpu setup 
### System Info

- `transformers` version: 4.26.0.dev0
- Platform: Linux-4.19.0-23-cloud-amd64-x86_64-with-glibc2.28
- Python version: 3.9.13
- Huggingface_hub version: 0.10.1
- PyTorch version (GPU?): 1.11.0+cu113 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes


### Who can help?

@sgguger

### Information

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

### Tasks

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

### Reproduction

When running on a multi-gpu setup, the output of `trainer.predict()` misses samples even if the parameter `dataloader_drop_last` is set to `False` in the `TrainingArguments`. 

[Here is an example colab notebook](https://drive.google.com/file/d/1pvnDkAhkMjLkUZVvindsFqhA_DRLveUk/view?usp=sharing) that could be run on a multi-gpu instance to reproduce the issue.

As an example, I had 2 gpus on my machine, my dataset had a size of 25000 samples, batch size per device was set to 256 (then 512 in total). When calling `trainer.predict()`, the output I got had a size of 24951.

The number of missing samples (25000-24951= 49) corresponds to the size of the dataloader (output of `len(trainer.get_test_dataloader()`)

### Expected behavior

Output of `trainer.predict(dataset)` should be the same length as dataset
The model used in this notebook (`Roberta0`) does not follow the requirements of the `Trainer` as described [here](https://huggingface.co/docs/transformers/main_classes/trainer) (see the big square in red). The output of the model should be a tuple, a dictionary or a `ModelOutput`, but it can't be a simple tensor. This is the reason for the problem seen.<|||||>Thanks for your quick response @sgugger 