def train_step(loss, *cache):
      """Training step function for use inside a while loop."""
      del loss  # unused; required in function signature.
      inputs = dequeue_fn()
      features, labels = inputs.features_and_labels()

      # Consume the current cache
      estimator_spec = self._verify_estimator_spec(
          self._call_model_fn(features, labels, cache=cache))

      # Retrieve the new returned cache
      """
        `cache` consists of a list of tensors, potentially empty (of length 0)
      """
      cache = estimator_spec.cache
      loss, train_op = estimator_spec.loss, estimator_spec.train_op

      if isinstance(estimator_spec, model_fn_lib._TPUEstimatorSpec):  # pylint: disable=protected-access
        captured_scaffold_fn.capture(estimator_spec.scaffold_fn)
      else:
        captured_scaffold_fn.capture(None)

      captured_training_hooks.capture(estimator_spec.training_hooks)

      tracing_ops = []
      if tensor_tracer.TensorTracer.is_enabled():
        tt = tensor_tracer.TensorTracer()
        loss, tracing_ops = tt.trace_tpu(ops.get_default_graph(), loss,
                                         self._ctx.num_replicas)

      # We must run train_op to update the variables prior to running the
      # outfeed.
      with ops.control_dependencies([train_op]+tracing_ops):
        host_call_outfeed_ops = []
        if (isinstance(estimator_spec, model_fn_lib._TPUEstimatorSpec)  # pylint: disable=protected-access
            and estimator_spec.host_call is not None):
          host_call.record({'host_call': estimator_spec.host_call})
          host_call_outfeed_ops = host_call.create_enqueue_op()
        with ops.control_dependencies(host_call_outfeed_ops):
          return [array_ops.identity(loss)] + cache
Training step function for use inside a while loop.
src/tpu_estimator.py