Be more clear about train_step and test_step

Be more clear that the train_step and test_step are also moved to the accelerator.  This was not clear before.  This makes it explicit, which should help people understanding TPU efficiency.

Author: MalcolmSlaney

site/en/guide/distributed_training.ipynb

@@ -508,7 +508,7 @@
         "Here's what you need to change in your code:\n",
         "\n",
         "1. Create an instance of the appropriate `tf.distribute.Strategy`.\n",
-        "2. Move the creation of Keras model, optimizer and metrics inside `strategy.scope`.\n",
+        "2. Move the creation of Keras model, optimizer and metrics inside `strategy.scope`. Thus the code in the model's run(), train_step(), and test_step() will all be distributed and executed on the accelerator(s).\n",
         "\n",
         "TensorFlow distribution strategies support all types of Keras models—[Sequential](/keras/sequential_model.ipynb), [Functional](/keras/functional.ipynb), and [subclassed](/keras/custom_layers_and_models.ipynb).\n",
         "\n",