def METHOD_NAME(self, batch):
    logits_pos = self.core(batch["ids_pos"], batch["len_pos"], return_logits=True)
    logits_neg = self.core(batch["ids_neg"], batch["len_neg"], return_logits=True)
    # softmax to get the `probability` to rank pos/neg correctly
    return torch.exp(logits_pos) / (torch.exp(logits_pos) + torch.exp(logits_neg))