ValueError in predict function for ClassificationModel
## 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.6.1
- Platform: Linux-4.15.0-122-generic-x86_64-with-debian-buster-sid
- Python version: 3.7.2
- PyTorch version (GPU?): 1.8.1+cu102 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.
@LysandreJik 


## Information

Model I am using (Bert, XLNet ...): bert

The problem arises when using:
* [ ] my own modified scripts: 
 I noticed then when I decrease my train batch size from 32 to 16,  i get the following bug:
 
 
![image](https://user-images.githubusercontent.com/47935154/121743373-41728f80-cacf-11eb-9dd0-5c07fdba9010.png)

 (Please note that training for 10 epochs happens successfully)

The tasks I am working on is:
* [ ] my own task or dataset: 
My own dataset for binary classification of text documents.


## Expected behavior
Evaluation should happen as expected. I am not sure what to fix/how to investigate. Could not find much about it online. 
ValueError: could not broadcast input array from shape (16,2) into shape (4,2)
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.