fix(ml): some tests were failing because the train_test_split changes

fix(ml): some tests were failing because the train_test_split changes #33