title: Keras, how to retrieve layers outputs in a custom training loop ? (more precisely in the *train_step* function)
body: <p>In the Keras <a href="https://keras.io/guides/customizing_what_happens_in_fit/" rel="nofollow noreferrer">documentation</a> we see that it's possible to customize what happens during training after every batch, for instance if we want to compute our own loss, we can override the <em>train_step</em> function:</p>
  <pre><code>class CustomModel(keras.Model):
      def train_step(self, data):
          x, y = data
  
          with tf.GradientTape() as tape:
              y_pred = self(x, training=True)  # Forward pass
              # Compute our own loss
              loss = keras.losses.mean_squared_error(y, y_pred)
  </code></pre>
  <p>But, what if I want to use some outputs of some layers in order to compute a custom loss after every batch ?</p>
  <p>Is there a way to have access to some layers outputs in the <em>train_step</em> function ?</p>