trainTest function finally working
More changes to the traintest script