Currently train, val and test split is monolithic. See data_preprocess.py -> make_train_val_test_split()
If the split is so that class percentages are not maintained then thats not a good way to split the dataset.
First, we need to check if class percentages are maintained in the splits or not?
If not, then split the dataset into train, val and test so that the class percentages match those of the original full dataset.