Train/Val Split with random_split | AlgoML Prep