import tensorflow as tf

      input_mask=inp_mask)
  summary = xlnet_model.get_pooled_out(FLAGS.summary_type, FLAGS.use_summ_proj)

  with tf.variable_scope("logits"):
    logits = tf.layers.dense(summary, 1,
        kernel_initializer=xlnet_model.get_initializer())
    logits = tf.reshape(logits, [bsz_per_core, 4])

    one_hot_target = tf.one_hot(label, 4)
    per_example_loss = -tf.reduce_sum(
        tf.nn.log_softmax(logits) * one_hot_target, -1)
    total_loss = tf.reduce_mean(per_example_loss)

  return total_loss, per_example_loss, logits
