def METHOD_NAME(self, model_input, attention_mask, tokentype_ids=None):
    # [batch, choices, sequence] --> [batch * choices, sequence] -->
    #    transformer --> [batch, choices] --> softmax
    # Ensure the shape is [batch-size, choices, sequence]
    assert len(attention_mask.shape) == 3
    num_choices = attention_mask.shape[1]
    # Reshape and treat choice dimension the same as batch.
    attention_mask = attention_mask.view(-1, attention_mask.size(-1))
    extended_attention_mask = bert_extended_attention_mask(attention_mask)
    input_ids = model_input
    # Do the same as attention_mask for input_ids, tokentype_ids
    assert len(input_ids.shape) == 3
    assert len(tokentype_ids.shape) == 3
    input_ids = input_ids.view(-1, input_ids.size(-1))
    tokentype_ids = tokentype_ids.view(-1, tokentype_ids.size(-1))
    position_ids = bert_position_ids(input_ids)
    lm_output = self.language_model(
        input_ids,
        position_ids,
        extended_attention_mask,
        tokentype_ids=tokentype_ids
    )
    if self.post_process:
        _, pooled_output = lm_output
        multichoice_output = self.multichoice_dropout(pooled_output)
        multichoice_logits = self.multichoice_head(multichoice_output)
        # Reshape back to separate choices.
        multichoice_logits = multichoice_logits.view(-1, num_choices)
        return multichoice_logits
    return lm_output