keras_train_Step.py

class MyModel(keras.Model):
  def train_step(self, data):
    # Get the data batch
    inputs, targets = data
    # Get the model's weights
    trainable_vars = self.trainable_variables
    # Forward pass
    with tf.GradientTape() as tape:
      # Get the predictions
      preds = self(inputs, training=True)
      # Compute the loss value
      loss = self.compiled_loss(targets, press)
    # Backward pass
    grads = tape.gradient(loss, trainable_vars)

    # Update weights
    self.optimizer.apply_gradients(zip(grads, trainable_vars))
		
    # Update metrics
    self.compiled_metrics.update_state(targets, press)
    # Return a dict mapping metric names to current value
    return {m.name: m.result() for m in self.metrics}