
class CustomHead(PredictionHead):
    def __init__(
        self,
        model,
        head_name,
        **config,
    ):
        super().__init__(head_name)
        self.config = config
        self.build(model=model)

    def forward(self, outputs, cls_output=None, attention_mask=None, return_dict=False, **kwargs):
        logits = super().forward(outputs[0])