Trainer returns logits of only one sequence instead of entire evaluation dataset
## 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  (also reproduced the same issue with 3.5.1)
- Platform: Google Colab
- Python version: 3.6
- PyTorch version (GPU?): 1.7.0+cu101
- Tensorflow version (GPU?): no
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: (don't know. Probably not)

### Who can help
@sgugger
(The use-case is basically the same as reported in this issue: https://github.com/huggingface/transformers/issues/9160
But I'm opening a new issue, because this is a separate problem) 

## Information

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

The tasks I am working on is:
* [x] my own task or dataset: (give details below)

## To reproduce

**Description:**
I’m trying to fine-tune a pre-trained NLI model (`ynie/roberta-large-snli_mnli_fever_anli_R1_R2_R3-nli`) on a dataset of around 276.000 hypothesis-premise pairs. I’m following the instructions from the docs [here](https://huggingface.co/transformers/custom_datasets.html) and [here](https://huggingface.co/transformers/training.html). 
- When I run the training, it seems like the fine-tuning works (it does the training and saves the checkpoints). 
- But during evaluation, the trainer/model seems to only compute the logits for one sequence instead of the entire training sequence. This means that the resulting metrics don't make sense. This also makes me doubt whether the training is actually executed on the entire training dataset, or only on the first sequence (the loss stays between 7.2 and 7.0 for most of the training steps).
- I've tried to change the input and dataset format in many different ways. One possible issue is that I am somehow only passing the same sequence into the trainer. But I've made sure that this is not the case (see also the printed input_ids in the code snippet below). 


```
### load model and tokenizer
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
max_length = 256
hg_model_hub_name = "ynie/roberta-large-snli_mnli_fever_anli_R1_R2_R3-nli"
tokenizer = AutoTokenizer.from_pretrained(hg_model_hub_name)
model = AutoModelForSequenceClassification.from_pretrained(hg_model_hub_name)  # num_labels=3
model.config

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device: {device}")
https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html
model.to(device)
model.train();

# ... some data preprocessing

### use hf datasets object + tokenize
from datasets import Dataset

# create hf dataset objects
dataset_train = Dataset.from_pandas(df_train)
dataset_val = Dataset.from_pandas(df_val)
dataset_test = Dataset.from_pandas(df_test)

## tokenize on dataset object
def tokenize(batch):
    return tokenizer(batch['premise'], batch['hypothesis'], max_length=max_length, return_token_type_ids=True, truncation=False, padding=True)
dataset_train = dataset_train.map(tokenize, batched=True, batch_size=len(df_train))
dataset_val = dataset_val.map(tokenize, batched=True, batch_size=len(df_val))
dataset_test = dataset_test.map(tokenize, batched=True, batch_size=len(df_test))

# to tensors
dataset_train.set_format('torch', columns=['input_ids', 'attention_mask', 'label', 'token_type_ids']) # format_kwargs=torch.LongTensor()
dataset_val.set_format('torch', columns=['input_ids', 'attention_mask', 'label', 'token_type_ids'])
dataset_test.set_format('torch', columns=['input_ids', 'attention_mask', 'label', 'token_type_ids'])

print(dataset_val["input_ids"][:2])
# ! to show that the tokenized input sequences are not the same
# output: 
tensor([[    0, 45784,     5,  3302,     9,  5222,   274,  1729,     7,  1719,
             5,  1403,    12, 28326,     7,  1157,  1232,     6,    11,  1989,
         10801,  1041,     6,  2388,  8243,    11,  1818,   173,     4,     2,
             2,   133,  2788,    16,    59,   592,  1134,     4,   286,  1246,
            35,  6300,     8,  3111,     6,   786,    12, 12063, 15343,  1134,
             6,  6610,  1134,     6,  1692,  1380,     8,  2038,  1134,     6,
           223, 25943, 33421,  5688,  1134,     2,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1],
        [    0,  1121,   570,    51,  1381,     7,   989,   106,    19,  1085,
            13,    10,   353,     4,     2,     2,   133,  2788,    16,    59,
          6642,     8,  1318,     9,   301,     4,   286,  1246,    35,  3039,
          2591,     6,  1265, 22830,     6,  2040,     6,  6642,   194, 22830,
             6,  6642,   194,  2919,     6,  9057,     6,  1265,  2919,     2,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
             1,     1,     1,     1,     1,     1,     1]])

# compute metrics with trainer https://colab.research.google.com/drive/1-JIJlao4dI-Ilww_NnTc0rxtp-ymgDgM?usp=sharing#scrollTo=N8J-TLhBuaOf
from sklearn.metrics import accuracy_score, precision_recall_fscore_support

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    print(pred.predictions, pred.predictions.argmax(-1))
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='binary', pos_label=0)
    acc = accuracy_score(labels, preds)
    return {
        'accuracy': acc,
        'f1': f1,
        'precision': precision,
        'recall': recall
    }

## training
from transformers import Trainer, TrainingArguments

# https://huggingface.co/transformers/main_classes/trainer.html#transformers.TrainingArguments
training_args = TrainingArguments(
    output_dir='./results',          # output directory
    num_train_epochs=1,              # total number of training epochs
    per_device_train_batch_size=16,   # batch size per device during training
    per_device_eval_batch_size=16,    # 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
    logging_steps=100,
    fp16=True
)

trainer = Trainer(
    model=model,                         # the instantiated 🤗 Transformers model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=dataset_train,         # training dataset
    eval_dataset=dataset_val,            # evaluation dataset
    compute_metrics=compute_metrics,
    tokenizer=tokenizer
)

trainer.train()

### evaluate
# for some reason the logits are the the same for each prediction
trainer.evaluate()

[[ 2.055 -7.992  2.053]
 [ 2.055 -7.992  2.053]
 [ 2.053 -7.984  2.05 ]
 ...
 [ 2.055 -7.992  2.053]
 [ 2.055 -7.992  2.053]
 [ 2.055 -7.992  2.053]] [0 0 0 ... 0 0 0]

{'eval_accuracy': 0.499,
 'eval_f1': 0.6657771847898598,
 'eval_loss': 0.6932016611099243,
 'eval_precision': 0.499,
 'eval_recall': 1.0}
```


<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

Get different logits for each sequence/model prediction in therefore get meaningful metrics output on the evaluation dataset.

<!-- A clear and concise description of what you would expect to happen. -->

**Update:** I slightly changed my script based on the [text_classification example notebook](https://github.com/huggingface/notebooks/blob/master/examples/text_classification.ipynb) and now it works. Don't know what made the difference, but in case someone else has a similar problem, I recommend following the steps in this script. 