import tensorflow as tf

      end_top_log_probs, end_top_index = tf.nn.top_k(
          end_log_probs, k=FLAGS.end_n_top)
      end_top_log_probs = tf.reshape(
          end_top_log_probs,
          [-1, FLAGS.start_n_top * FLAGS.end_n_top])
      end_top_index = tf.reshape(
          end_top_index,
          [-1, FLAGS.start_n_top * FLAGS.end_n_top])

  if is_training:
