Bug fix for apply_chunking_to_forward chunking dimension check

Chunking should be in the chunking dimension, an exception was raised if the complete shape of the inputs was not the same rather than only the chunking dimension

# What does this PR do?
Fixes #8349


## 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?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to the 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?

@patrickvonplaten Incorporating the fix from the issue at https://github.com/huggingface/transformers/issues/8349
Let me know if you want a unit test or something! Ran all the tests too:
<img width="936" alt="Screen Shot 2020-11-07 at 3 29 26 PM" src="https://user-images.githubusercontent.com/5157240/98450932-1e35cb00-210f-11eb-855a-2457c44e243b.png">


