Solved by ptrblck in post #2 
              
              
                You could use sklearn.model_selection.train_test_split with the stratify argument to create the training and testing indices.