import tensorflow as tf

    epsilon = tf.constant(value=1e-7)
    labels = tf.to_float(labels)
    # labels = tf.to_float(tf.reshape(labels, (-1, num_classes)))

    softmax = tf.nn.softmax(logits) + epsilon

    if head is not None:
      cross_entropy = -tf.reduce_sum(tf.mul(labels * tf.log(softmax), head), axis=[1])
