How to split test and train data keeping equal proportions of each class?

[ptrblck]

If you load the dataset completely before passing it to the Dataset and DataLoader classes, you could use scikit-learn’s train_test_split with the stratified option.