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



