from tensorflow.python.ops import math_ops

  def get_weight_tensor(self, features):
    if not self._weight_column_name:
      return None
    else:
      return array_ops.reshape(
          math_ops.to_float(features[self._weight_column_name]),
          shape=(-1,))

  def loss(self, logits, target, features):
    """Returns loss tensor for this head.
