kfold.py

from sklearn.model_selection import KFold

kfold = KFold(n_splits=10)

for i, (trn_idx, val_idx) in enumerate(kfold.split(X_train)):
    print("Working on {}th Fold".format(i))
    
    train_multi_ds = TensorDataset(X_train[trn_idx], y_multi_train[trn_idx])
    val_multi_ds = TensorDataset(X_train[val_idx], y_multi_train[val_idx])
    
    train_multi_loader = DataLoader(train_multi_ds, batch_size=128, shuffle=True)
    val_multi_loader = DataLoader(val_multi_ds, batch_size=128, shuffle=True)