Merge pull request #1787 from peastman/training

    @tf.function(experimental_relax_shapes=True)

    def apply_gradient_for_batch(inputs, labels, weights, loss):

      with tf.GradientTape() as tape:

        outputs = self.model(inputs)

        outputs = self.model(inputs, training=True)

        if isinstance(outputs, tf.Tensor):

          outputs = [outputs]

        if self._loss_outputs is not None:

  raise ValueError("Incompatible shapes for outputs and labels: %s versus %s" %

                   (str(shape1), str(shape2)))

def _ensure_float(output, labels):

  """Make sure the outputs and labels are both floating point types."""

  if output.dtype not in (tf.float32, tf.float64):