def METHOD_NAME(self,
            input_ids=None,
            attention_mask=None,
            token_type_ids=None,
            position_ids=None,
            head_mask=None,
            inputs_embeds=None,
            output_attentions=None,
            output_hidden_states=None,
            return_dict=None,
            **kwargs):
    """
    Returns task head applied to language model outputs and any additional arguments supplied via `kwargs`.
    See HF transformers documentation for more details on other parameters.
    """
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.model_(
        input_ids,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
        position_ids=position_ids,
        head_mask=head_mask,
        inputs_embeds=inputs_embeds,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )
    return self.task_head(outputs, input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, **kwargs)