trainer.evaluate() expects batch_size to match target batch_size
@LysandreJik
@sgugger

## 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.0.1
- Platform: Windows/Ubuntu 18.04.3
- Python version: 3.7.6
- PyTorch version (GPU?): 1.7.1 CPU
- Using distributed or parallel set-up in script?: Nope


## Information

Model I am using ('deepset/gbert-base'):

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)


The problem I get is the following when I call the trainer.evaluate() function:

```Bash
Traceback (most recent call last):
  File "fine_tune_bert.py", line 174, in <module>
    trainer.evaluate()
  File "/home/rouven/anaconda3/lib/python3.7/site-packages/transformers/trainer.py", line 1259, in evaluate
    ignore_keys=ignore_keys,
  File "/home/rouven/anaconda3/lib/python3.7/site-packages/transformers/trainer.py", line 1363, in prediction_loop
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "/home/rouven/anaconda3/lib/python3.7/site-packages/transformers/trainer.py", line 1469, in prediction_step
    outputs = model(**inputs)
  File "/home/rouven/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/rouven/anaconda3/lib/python3.7/site-packages/transformers/models/bert/modeling_bert.py", line 1363, in forward
    loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
  File "/home/rouven/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/rouven/anaconda3/lib/python3.7/site-packages/torch/nn/modules/loss.py", line 962, in forward
    ignore_index=self.ignore_index, reduction=self.reduction)
  File "/home/rouven/anaconda3/lib/python3.7/site-packages/torch/nn/functional.py", line 2468, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/home/rouven/anaconda3/lib/python3.7/site-packages/torch/nn/functional.py", line 2262, in nll_loss
    .format(input.size(0), target.size(0)))
ValueError: Expected input batch_size (18) to match target batch_size (6).

```

I'm doing a multiclass classification problem. With six classes, that is why i'm replacing the classifyer here.

```Python
model = BertForSequenceClassification.from_pretrained('deepset/gbert-base', proxies=charite_proxy)
model.classifier = torch.nn.Linear(768, 6)
```


I had the same problem with the trainer.train() call before overwriting the compute_loss function. Which looks like this now:


```Python
class MultilabelTrainer(Trainer):

    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs[0]

        global weigths
        global lambda_reg
        reg_lambda = lambda_reg

        weight = weights
        criterior = CrossEntropyLoss(weight=weight.to(device))

        loss = criterior(logits, labels)
        loss += calculate_l2_reg(model, reg_lambda)

        return (loss, outputs) if return_outputs else loss
```


Further my training setup looks like this:

```Python
EPOCHS = 3
LEARNING_RATE = 2e-5 
BATCH_SIZE = 32

training_args = TrainingArguments(
    output_dir='./results',          # output directory
    num_train_epochs=EPOCHS,         # total # of training epochs
    per_device_train_batch_size=BATCH_SIZE,  # batch size per device during training
    per_device_eval_batch_size=BATCH_SIZE,   # batch size for evaluation
    warmup_steps=500,                # number of warmup steps for learning rate scheduler
    weight_decay=0.01,               # strength of weight decay
    logging_dir='./logs',            # directory for storing logs
    no_cuda = True,
    seed = seed,
    learning_rate = LEARNING_RATE
)

model.train()

trainer = MultilabelTrainer(
    model=model,                         # the instantiated 🤗 Transformers model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=dataset,               # training dataset
    eval_dataset=test_dataset            # evaluation dataset  
)

trainer.train()
trainer.evaluate()
```

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

I'm not quite sure what you would need to know, but it is a dataset consisting of ~60k examples with 1 of 6 possible labels.


## Expected behavior
The expected behavior would be to get the evaluation metrics from the trainer.evaluate() call.

Hope you can help me.

Cheers
Rouven

You will need to update to the last version of Transformers (I'm seeing 4.0.1 in your report), we fixed this issue so the evaluation loop uses the `compute_loss` function too.<|||||>> You will need to update to the last version of Transformers (I'm seeing 4.0.1 in your report), we fixed this issue so the evaluation loop uses the `compute_loss` function too.

Thanks sgugger!

I also found this out after taking a dive into your code base. I overwrote the prediction_step function in my case, since i dont know if the rest of my code supports transformers 4.5.0.

You can close the issue now! :)