corrected train and train_single