It's probably less of an issue on "biggish" datasets, but a param for specifying stratification in random_split (which was merged here: #4435) would be nice. If the targets are class labels this could be done based on the integer labels. For regression problems, there are some interesting ideas posted here: http://scottclowe.com/2016-03-19-stratified-regression-partitions/

If this is of interest, we could discuss how to implement that efficiently, I am happy to help.

cc @ssnl @VitalyFedyunin