transformers

Adding new train_step logic to make things less confusing for users

#15994

Merged