def get_num_labels(labels_or_logits):
  """Returns the number of labels inferred from labels_or_logits."""
  if labels_or_logits.get_shape().ndims <= 1:
    return 1
  return labels_or_logits.get_shape()[1].value