Function train_step

tensorflow/python/keras/custom_training_loop_test.py:76–83

(x)

                                              input_shape=(10,))

  def train_step(x):
    with backprop.GradientTape() as tape:
      model(x)
      assert len(model.losses) == 2
      loss = math_ops.reduce_sum(model.losses)
    gradients = tape.gradient(loss, model.trainable_weights)
    optimizer.apply_gradients(zip(gradients, model.trainable_weights))
    return loss

  if defun:
    train_step = def_function.function(train_step)