from tensorflow.python.ops import math_ops

      return self._dnn_logits(features)
    else:
      return self._linear_logits(features)

  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, weight_tensor):
    if self._n_classes < 2:
      loss_vec = math_ops.square(logits - math_ops.to_float(target))
    elif self._n_classes == 2:
      loss_vec = nn.sigmoid_cross_entropy_with_logits(logits,
                                                      math_ops.to_float(target))
