import tensorflow as tf

    return total_loss, per_example_loss, logits


def get_qa_outputs(FLAGS, features, is_training):
  """Loss for downstream span-extraction QA tasks such as SQuAD."""

  inp = tf.transpose(features["input_ids"], [1, 0])
  seg_id = tf.transpose(features["segment_ids"], [1, 0])
  inp_mask = tf.transpose(features["input_mask"], [1, 0])
  cls_index = tf.reshape(features["cls_index"], [-1])

  seq_len = tf.shape(inp)[0]
