Error while training fb-bart-mnli

I am trying to train the bove model using the example given in the notebook.

The only difference between this note book and mine are as follows:

Instead of loading the data set from the hub I am loading it from my local storage on Sagemaker notebook instance.

train_dataset = load_dataset("csv", data_files="./data/prompt_data.csv")
test_dataset = load_dataset("csv", data_files="./data/prompt_test.csv")

instead of :

train_dataset, test_dataset = load_dataset('imdb', split=['train', 'test'])

rest of the notebook is exactly the same.

When i run the training I get th following error:
2023-06-04 04:15:58,040 sagemaker-training-toolkit INFO Waiting for the process to finish and give a return code.
2023-06-04 04:15:58,040 sagemaker-training-toolkit INFO Done waiting for a return code. Received 1 from exiting process.
2023-06-04 04:15:58,041 sagemaker-training-toolkit ERROR Reporting training FAILURE
2023-06-04 04:15:58,041 sagemaker-training-toolkit ERROR ExecuteUserScriptError:
ExitCode 1
ErrorMessage “│ ❱ 64 │ │ │ raise KeyError( │
│ 65 │ │ │ │ f"Invalid key: {k}. Please first select a split. For │
│ 66 │ │ │ │ f"my_dataset_dictionary['{suggested_split}'][{k}]. │
│ 67 │ │ │ │ f"Available splits: {sorted(self)}” │
╰──────────────────────────────────────────────────────────────────────────────╯
KeyError: “Invalid key: 0. Please first select a split. For example
my_dataset_dictionary['train'][0]. Available splits: [‘train’]”
0%| | 0/1 [00:00<?, ?it/s]"
Command “/opt/conda/bin/python3.9 train.py --epochs 1 --model_name facebook/bart-large-mnli --train_batch_size 32”
2023-06-04 04:15:58,041 sagemaker-training-toolkit ERROR Encountered exit_code