minor cleaning in training code