Python KeyError: 0

I'm using pytorch to train UNet on 3D segmentation. In my kickstart script, I tried to create3 keys from my data and name them: "train", "val", "test". I'm getting an error in the "test" dictionary called KeyError: 0. I'm not sure why "train" and "val" are running fine but just "test" is giving the error. My data are 3D Nifti files.

import os
import json

from experiments.UNetExperiment import UNetExperiment
from data_prep.HippocampusDatasetLoader import LoadHippocampusData


class Config:
    def __init__(self):
        self.name = "hippocampus_unet"
        self.root_dir = r"/content/drive/My Drive/"
        self.n_epochs = 10
        self.learning_rate = 0.0002
        self.batch_size = 8
        self.patch_size = 64
        self.test_results_dir = "/content/drive/My Drive/"

if __name__ == "__main__":
    c = Config()
    print("Loading data...")
    data = LoadHippocampusData(c.root_dir, y_shape = c.patch_size, z_shape = c.patch_size)
    
    keys = range(len(data))
    split = dict()

    np.random.seed(42)

    length = len(data)
    new_index = np.random.permutation(length)

    train_val_index = new_index[:int(length*0.8)]
    train_index = train_val_index[:int(length*0.8)]
    val_index = list(set(train_val_index) - set(train_index))
    test_index = list(set(new_index) - set(train_val_index))

    split["train"] = {key: [value[i] for i in train_index] for key, value in split.items()}
    split["val"] = {key: [value[i] for i in val_index] for key, value in split.items()}
    split["test"] = {key: [value[i] for i in test_index] for key, value in split.items()}

    exp = UNetExperiment(c, split, data)
    exp.run()
    results_json = exp.run_test()
    results_json["config"] = vars(c)

    with open(os.path.join(exp.out_dir, "results.json"), 'w') as out_file:
        json.dump(results_json, out_file, indent=2, separators=(',', ': '))

Then I get KeyError 0 at split["test"] line.

enter image description here