In the body of the train_step method, we implement a regular training update, similar to what you are already familiar with. Importantly, we compute the loss via self.compiled_loss, which wraps the loss(es) function(s) that were passed to compile().