text: Title: IndexError: Invalid key: 0 is out of bounds for size 0
  Body: ### Describe the bug
  
  I am trying to fine-tune llama2-7b model in GCP. The notebook I am using for this can be found [here](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_llama2_peft_finetuning.ipynb).  
  
  When I use the dataset given in the example, the training gets successfully completed (example dataset can be found [here](https://huggingface.co/datasets/timdettmers/openassistant-guanaco)).  
  However when I use my own dataset which is in the same format as the example dataset, I get the below error (my dataset can be found [here](https://huggingface.co/datasets/kk2491/finetune_dataset_002)).   
  
  ![image](https://github.com/huggingface/datasets/assets/38481564/47fa2de3-95e0-478b-a35f-58cbaf90427a)  
  
  I see the files are being read correctly from the logs:  
  
  ![image](https://github.com/huggingface/datasets/assets/38481564/b0b6316c-2cc7-476c-9674-ca2222c8f4e3)
  
  
  
  ### Steps to reproduce the bug
  
  1. Clone the [vertex-ai-samples](https://github.com/GoogleCloudPlatform/vertex-ai-samples) repository. 
  2. Run the [llama2-7b peft fine-tuning](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_llama2_peft_finetuning.ipynb).
  3. Change the dataset `kk2491/finetune_dataset_002`   
  
  ### Expected behavior
  
  The training should complete successfully, and model gets deployed to an endpoint. 
  
  
  
  ### Environment info
  
  Python version : Python 3.10.12
  Dataset : https://huggingface.co/datasets/kk2491/finetune_dataset_002  
  
  