import tensorflow as tf

            next_sentence_example_loss,
            next_sentence_log_probs,
        ) = get_next_sentence_output(
            bert_config, model.get_pooled_output(), next_sentence_labels
        )

        total_loss = masked_lm_loss + next_sentence_loss

        tvars = tf.trainable_variables()

        initialized_variable_names = {}
        scaffold_fn = None
        if init_checkpoint:
            (
                assignment_map,
                initialized_variable_names,
