Minor changes to SGDTrainer