import tensorflow as tf


  else:
    q = tf.nn.softmax(q_logits)
    p = tf.nn.softmax(p_logits)
    kl = tf.reduce_sum(q * (tf.log(q) - tf.log(p)), 1)

  num_labels = tf.reduce_sum(weights)
  num_labels = tf.where(tf.equal(num_labels, 0.), 1., num_labels)
