from tensorflow.python.ops import math_ops

    dnn_grads = grads[0:len(dnn_vars)]
    linear_grads = grads[len(dnn_vars):]

    train_ops = self._get_linear_training_ops(
        linear_grads, linear_vars) + self._get_dnn_training_ops(dnn_grads,
                                                                dnn_vars)

    train_step = control_flow_ops.group(*train_ops, name="combined_training_op")
    with ops.control_dependencies([train_step]):
      with ops.get_default_graph().colocate_with(global_step):
        return state_ops.assign_add(global_step, 1).op, loss

  def _run_metrics(self, predictions, targets, metrics, weights):
    result = {}
    targets = math_ops.cast(targets, predictions.dtype)
    for name, metric in six.iteritems(metrics or {}):
      if "weights" in inspect.getargspec(metric)[0]:
        result[name] = metric(predictions, targets, weights=weights)
      else:
        result[name] = metric(predictions, targets)

    return result

  def _get_eval_ops(self, features, targets, metrics=None):
    """See base class."""
    logits = self._logits(features)
    result = {"loss": metrics_lib.streaming_mean(self._loss(
        logits, targets,
