from tensorflow.python.ops import array_ops


  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))
    else:
      loss_vec = nn.sparse_softmax_cross_entropy_with_logits(
          logits, array_ops.reshape(target, [-1]))

    if weight_tensor is None:
      return math_ops.reduce_mean(loss_vec, name="loss")
    else:
      loss_vec = array_ops.reshape(loss_vec, shape=(-1,))
      loss_vec = math_ops.mul(
          loss_vec, array_ops.reshape(weight_tensor, shape=(-1,)))
      return math_ops.div(
