Test Train Validation Split Python 3 Module