Fixes to training and test program