key error when use trainer to fine_tuning a 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.1.1
- Platform: Linux-3.10.0-693.5.2.el7.x86_64-x86_64-with-centos-7.4.1708-Core
- Python version: 3.7.9
- PyTorch version (GPU?): 1.7.1 (False)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### 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.

 albert, bert, GPT2, XLM: @LysandreJik
 tokenizers: @mfuntowicz
 Trainer: @sgugger
 Speed and Memory Benchmarks: @patrickvonplaten
 Model Cards: @julien-c
 TextGeneration: @TevenLeScao
 examples/distillation: @VictorSanh
 nlp datasets: [different repo](https://github.com/huggingface/nlp)
 rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
 Text Generation: @patrickvonplaten @TevenLeScao
 Blenderbot: @patrickvonplaten
 Bart: @patrickvonplaten
 Marian: @patrickvonplaten
 Pegasus: @patrickvonplaten
 mBART: @patrickvonplaten
 T5: @patrickvonplaten
 Longformer/Reformer: @patrickvonplaten
 TransfoXL/XLNet: @TevenLeScao
 RAG: @patrickvonplaten, @lhoestq
 FSMT: @stas00
 examples/seq2seq: @patil-suraj
 examples/bert-loses-patience: @JetRunner
 ray/raytune: @richardliaw @amogkam
 tensorflow: @jplu
 examples/token-classification: @stefan-it
 documentation: @sgugger
 -->
@sgugger

## Information

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

The problem arises when using:
* [ ] the official example scripts: (give details below)
i am fine-tuning a text_claasifiction on dbpedia_14.and i followed this colab https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification.ipynb#scrollTo=TlqNaB8jIrJW


The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name) 
datset:dbpedia_14

## To reproduce

Steps to reproduce the behavior:

error
`File "train.py", line 69, in <module>
    trainer.train()
  File "/home/pliu3/projects/anaconda3/envs/calibration/lib/python3.7/site-packages/transformers/trainer.py", line 784, in train
    for step, inputs in enumerate(epoch_iterator):
  File "/home/pliu3/projects/anaconda3/envs/calibration/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 435, in __next__
    data = self._next_data()
  File "/home/pliu3/projects/anaconda3/envs/calibration/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 475, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/pliu3/projects/anaconda3/envs/calibration/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 44, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/pliu3/projects/anaconda3/envs/calibration/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 44, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
KeyError: 2`

code
```python
dataset_name = 'sem_eval_2014_task_1'

num_labels_size = 3

batch_size = 4

model_checkpoint = 'bert-base-uncased'

number_train_epoch = 5



def tokenize(batch):
    return tokenizer(batch['premise'], batch['hypothesis'],  truncation=True, )


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


model = BertForSequenceClassification.from_pretrained(model_checkpoint, num_labels=num_labels_size)
tokenizer = BertTokenizerFast.from_pretrained(model_checkpoint, use_fast=True)

train_dataset = load_dataset(dataset_name, split='train')
test_dataset = load_dataset(dataset_name, split='test')

train_encoded_dataset = train_dataset.map(tokenize, batched=True)
test_encoded_dataset = test_dataset.map(tokenize, batched=True)

args = TrainingArguments(
    output_dir='./results',
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    num_train_epochs=number_train_epoch,
    weight_decay=0.01,
    do_predict=True
)
trainer = Trainer(
    model=model,
    args=args,
    compute_metrics=compute_metrics,
    train_dataset=train_encoded_dataset,
    eval_dataset=test_encoded_dataset,
    tokenizer=tokenizer
)

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


Duplicate of #9636