Enhancement (Remove comments)
Fix train_return()