Fixed new training changes from last commit (was broken).