def METHOD_NAME(self, *args, **kwargs):
    kwargs['return_dict'] = True
    outputs = super(Model, self).METHOD_NAME(*args, **kwargs)
    return AttentionTokenClassificationModelOutput(
        loss=outputs.loss,
        logits=outputs.logits,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
    )