Rethinking the Training Loop

Training step

As already mentioned, the higher-order function that builds a training step function for us is taking the key elements of our training loop: model, loss, and optimizer. The actual training step function to be returned will have two arguments, namely, features and labels, and will return the corresponding loss value.

Creating the higher-order function for training step

Apart from returning the loss value, the inner perform_train_step() function below is the same as the code inside the loop in model training V0. The code should look like this: