def METHOD_NAME(self, logits, relations, entities):
    pred_relations = []
    for i, pred_label in enumerate(logits.argmax(-1)):
        if pred_label != 1:
            continue
        rel = {}
        rel["head_id"] = relations["head"][i]
        rel["head"] = (entities["start"][rel["head_id"]], entities["end"][rel["head_id"]])
        rel["head_type"] = entities["label"][rel["head_id"]]
        rel["tail_id"] = relations["tail"][i]
        rel["tail"] = (entities["start"][rel["tail_id"]], entities["end"][rel["tail_id"]])
        rel["tail_type"] = entities["label"][rel["tail_id"]]
        rel["type"] = 1
        pred_relations.append(rel)
    return pred_relations