How to avoid covariate shift in python and distribute classes in each train and test phase?
You need to call train_test_split(X, y, test_size=0.33, stratify=y) assuming that y contains your class labels. Alternatively, you can call StratifiedShuffleSplit directly.