Slow training time when using train_test_split

### Describe the bug

I have a big data set, with 1 class only (for binary classification). The dataset is highly imbalanced.

I was doing a manual stratify train test split using some for loops and if else, it was taking a lot of time for doing that , but the CNN training time was about 20 minutes.

after that I started using train_test_split with Stratify, it loads that data in 1 minute, BUT the training time now takes 2 hours!

I'm loading the data using a costume data generator, also the rest of the code is the same, the only difference is the loading data.

this is for split data using train test split 
```python
nodules_csv = pd.read_csv("/cropped_nodules.csv")
train_nodules, val_nodules = train_test_split(nodules_csv, test_size=0.10, stratify=nodules_csv['state'], random_state=42)


train_data = (train_nodules['SN'].astype(str) + '.npy').tolist()
train_labels = train_nodules['state'].tolist()


val_data = (val_nodules['SN'].astype(str) + '.npy').tolist()
val_labels = val_nodules['state'].tolist()

# Print the sizes of the splits to verify
print("x_train = ", len(train_data))
print("y_train = ", len(train_labels))
print("x_val = ", len(val_data))
print("y_val = ", len(val_labels))

```

this is for splitting data manually
```python
nodules_csv = pandas.read_csv("/cropped_nodules.csv")
base_dir = "/cropped_nodules/"
all_image_paths = os.listdir(base_dir)
all_image_paths = sorted(all_image_paths,key=lambda x: int(os.path.splitext(x)[0]))
abnormal_nodules= nodules.loc[nodules['state'] == 1]
normal_nodules= nodules.loc[nodules['state'] == 0]

y = pandas.concat([abnormal_nodules.sample(frac=0.10), normal_nodules.sample(frac=0.10)])
h= y['ID'].tolist()
h_string= map(str, h)
val_data = [item + '.npy' for item in h_string]

train_image_paths =[item for item in all_image_paths if item not in val_data]

val_label = []
train_labels =[]

for index in range(len(nodules)):
    if str(index) +'.npy' not in val_data:
        train_labels.append(nodules['state'].iloc[index])

for num_nodule in val_data:
    val_label.append(nodules.iloc[int(num_nodule[:-4]),1])

print("x_val = ", len(val_data))
print("x_train = ", len(train_image_paths))
print("y_val = ", len(val_label))

```
### Steps/Code to Reproduce

```python
nodules_csv = pd.read_csv("/cropped_nodules.csv")
train_nodules, val_nodules = train_test_split(nodules_csv, test_size=0.10, stratify=nodules_csv['state'], random_state=42)


train_data = (train_nodules['SN'].astype(str) + '.npy').tolist()
train_labels = train_nodules['state'].tolist()


val_data = (val_nodules['SN'].astype(str) + '.npy').tolist()
val_labels = val_nodules['state'].tolist()

# Print the sizes of the splits to verify
print("x_train = ", len(train_data))
print("y_train = ", len(train_labels))
print("x_val = ", len(val_data))
print("y_val = ", len(val_labels))
```

### Expected Results

No error is thrown

### Actual Results

No specific result

### Versions

```shell
System:
    python: 3.9.16 (main, Dec  7 2022, 01:11:58)  [GCC 7.5.0]
executable: /usr/bin/python3.9
   machine: Linux-5.4.0-150-generic-x86_64-with-glibc2.27
Python dependencies:
          pip: 22.2.2
   setuptools: 56.0.0
      sklearn: 0.24.2
        numpy: 1.23.1
        scipy: 1.9.0
       Cython: None
       pandas: 1.2.4
   matplotlib: 3.4.2
       joblib: 1.0.1
threadpoolctl: 2.1.0
Built with OpenMP: True
```

COMMENT:
I **guess** the problem is:

My manual code is sorting the dataset (they are physically closer to each other on the HDD).
The train_test_split takes the data randomly (they are far apart from each other) and that makes it slower to be loaded during the training

Labels: Question