We plan to add a way to define additional splits that just train and test in train_test_split.
For now you’d have to use it twice as you mentioned (or use a combination of Dataset.shuffle and Dataset.shard/select).
See the issue about extending train_test_split here 22