Fix exception in prediction loop occurring for certain batch sizes
# What does this PR do?

Fixes #12349 

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@sgugger 

Thanks, I can't push to your branch, so could you push an empty commit since circleCI decided to not run on your PR?
```
git commit --allow-empty -m "Trigger CI"
```
and then a push.<|||||>Ok, there are lots of failures that seem related to the PyTorch release. Could you rebase on master (I really want to make sure this does not break any example before merging)? Thank you.<|||||>> Ok, there are lots of failures that seem related to the PyTorch release. Could you rebase on master (I really want to make sure this does not break any example before merging)? Thank you.

Done, I must have based this off from a fork that was a few weeks old....<|||||>I'll be off now for the weekend, but the box "Allow edits by maintainers" is checked, so feel free to adapt as necessary.<|||||>We're all good, thanks for walking through this with me!