def METHOD_NAME(self, input_ids, token_type_ids=None, attention_mask=None):
    # Only return lm_logits
    return self.model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)