IndexError: Invalid key: 14 is out of bounds for size 0
Hi ! Could you check the length of the `self.dataset` object (i.e. the Dataset object passed to the data loader) ? It looks like the dataset is empty.
Not sure why the SWA optimizer would cause this though. Any updates on this? 
The same error occurred to me too when running `cardiffnlp/twitter-roberta-base-sentiment` on a custom dataset. This happened when I tried to do `model = torch.nn.DataParallel(model, device_ids=[0, 1, 2, 3])` without using sagemaker distribution. 
Python: 3.6.13
datasets: 1.6.2 Hi @ruisi-su, do you have this issue while using SWA as well, or just data parallel ?

If you have a code example to reproduce this issue it would also be helpful @lhoestq I had this issue without SWA. I followed [this](https://github.com/huggingface/notebooks/blob/master/sagemaker/03_distributed_training_data_parallelism/sagemaker-notebook.ipynb) notebook to utilize multiple gpus on the [roberta-base](https://huggingface.co/cardiffnlp/twitter-roberta-base-sentiment) model. This tutorial could only work if I am on `ml.p3.16xlarge`, which I don't have access to. So I tried using just `model = torch.nn.DataParallel(model, device_ids=[0, 1, 2, 3]` before calling `trainer.fit()`. But maybe this is not the right way to do distributed training. I can provide a code example if that will be more helpful. It might be an issue with old versions of `datasets`, can you try updating `datasets` ? FYI I encountered the exact same error using the latest versions of `datasets`, `transformers` and `pyarrow`, without using any kind of SWA or dataparallel: 

```
# packages in environment at C:\Users\zhang\mambaforge:
#
# Name                    Version                   Build  Channel
cudatoolkit               11.0.3               h3f58a73_9    https://mirrors.ustc.edu.cn/anaconda/cloud/conda-forge
datasets                  1.17.0                   pypi_0    pypi
pyarrow                   6.0.1                    pypi_0    pypi
pytorch                   1.7.1           py3.9_cuda110_cudnn8_0    pytorch
tornado                   6.1              py39hb82d6ee_2    https://mirrors.ustc.edu.cn/anaconda/cloud/conda-forge
```

```
> python --version
> 3.9.7
``` Same error here! Datasets version `1.18.3` freshly updated.

`IndexError: Invalid key: 90 is out of bounds for size 0`

My task is finetuning the model for token classification.

**Solved**: I make a mistake while updating the dataset during the map, you should check that you return the correct values.
 cc @sgugger This probably comes from the `Trainer` removing all the columns of a dataset, do you think we can improve the error message in this case ? The `Trainer` clearly logs when it removes columns in the dataset. I'm not too sure of where the bug appears as I haven't seen a clear reproducer. Happy to display a more helpful error message, but I'd need a reproducer to see what the exact problem is to design the right test and warning :-)  Well, if I can try to suggest how to reproduce, please try by do not returning any updated content in the map function used to tokenize input (e.g., in TokenClassification). I can leave here my wrong version for reference:

```python
def preprocess_function(examples):

    text = examples["text"]
    
    inputs = tokenizer(
        text,
        max_length=512,
        truncation="only_second",
        return_offsets_mapping=True,
        padding="max_length",
    )

    offset_mapping = inputs.pop("offset_mapping")
    # ... processing code

    inputs["labels"] = label_ids
    #return inputs
    
train_ds = train_ds.map(preprocess_function, batched=False)
test_ds = test_ds.map(preprocess_function, batched=False)
eval_ds = eval_ds.map(preprocess_function, batched=False)
```

Of course, returning inputs solved the problem. As suggestion, a possible error message could display "IndexError: the `key` required by trainer are not found in the dataset" (just an hypothesis, I think there could be something better). 

Please tell me if you need more details to reproduce, glad to help! That's the thing though. The `Trainer` has no idea which inputs are required or not since all models can have different kinds of inputs, and it can work for models outside of the Transformers library. I can add a clear error message if I get an empty batch, as this is easy to detect, but that's pretty much it. I think that it could be enough to ease the identification of the problem. Done in [this commit](https://github.com/huggingface/transformers/commit/c87cfd653c4de3d4743a9ae09d749282d94d5829)
## Describe the bug
A clear and concise description of what the bug is.

Hi. I am trying to implement stochastic weighted averaging optimizer with transformer library as described here https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging/ , for this I am using a run_clm.py codes which is working fine before adding SWA optimizer, the moment I modify the model with `swa_model = AveragedModel(model)` in this script, I am getting the below error,  since I am NOT touching the dataloader part, I am confused why this is occurring, I very much appreciate your opinion on this @lhoestq 

## Steps to reproduce the bug
```
Traceback (most recent call last):
  File "run_clm.py", line 723, in <module>
    main()
  File "run_clm.py", line 669, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/user/dara/libs/anaconda3/envs/success/lib/python3.7/site-packages/transformers/trainer.py", line 1258, in train
    for step, inputs in enumerate(epoch_iterator):
  File "/user/dara/libs/anaconda3/envs/success/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 435, in __next__
    data = self._next_data()
  File "/user/dara/libs/anaconda3/envs/success/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 "/user/dara/libs/anaconda3/envs/success/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 "/user/dara/libs/anaconda3/envs/success/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]
  File "/user/dara/libs/anaconda3/envs/success/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 1530, in __getitem__
    format_kwargs=self._format_kwargs,
  File "/user/dara/libs/anaconda3/envs/success/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 1517, in _getitem
    pa_subtable = query_table(self._data, key, indices=self._indices if self._indices is not None else None)
  File "/user/dara/libs/anaconda3/envs/success/lib/python3.7/site-packages/datasets/formatting/formatting.py", line 368, in query_table
    _check_valid_index_key(key, size)
  File "/user/dara/libs/anaconda3/envs/success/lib/python3.7/site-packages/datasets/formatting/formatting.py", line 311, in _check_valid_index_key
    raise IndexError(f"Invalid key: {key} is out of bounds for size {size}")
IndexError: Invalid key: 14 is out of bounds for size 0
```



## Expected results
not getting the index error

## Actual results
Please see the above

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: datasets                  1.12.1 
- Platform: linux
- Python version: 3.7.11 
- PyArrow version:  5.0.0 



