Ideally we want to use now for 2.2 `train_step`: https://www.tensorflow.org/api_docs/python/tf/keras/Model?version=nightly#train_step
Ideally we want to use now for 2.2
train_step: https://www.tensorflow.org/api_docs/python/tf/keras/Model?version=nightly#train_step