def METHOD_NAME(self, input_id, mask, head_entity, tail_entity):
    _, pooled_output = self.bert(input_ids=input_id, attention_mask=mask, return_dict=False)
    nlp_output = self.bert_reduce_stack(pooled_output)
    score = self.pdist(head_entity + nlp_output, tail_entity)
    output = self.relu(score)
    return output