from sklearn.model_selection import StratifiedShuffleSplit sss=StratifiedShuffleSplit(n_splits=1,test_size=0.25,random_state=0) for train_index, test_index in sss.split(df, df["km_driven_cat"]): train_set = df.iloc[train_index] test_set = df.iloc[test_index]