关于头尾实体特征和关系特征的问题 about pfn HOT 5 CLOSED

    ner_mask = torch.sum(ner_pred, dim=1).float()
    ner_mask = torch.sum(ner_mask, dim=-1).float()
    ner_mask = torch.where(ner_mask > 0, torch.ones_like(ner_mask), torch.zeros_like(ner_mask))

    seq_len = ner_mask.size(0)
    e1 = ner_mask.unsqueeze(0).repeat(seq_len, 1, 1)
    e2 = ner_mask.unsqueeze(1).repeat(1, seq_len, 1)
    ner_mask = e1 * e2

    ner_mask = ner_mask.unsqueeze(-1).repeat(1, 1, 1, len(self.rel2idx))
    complete_re_pred = re_pred * ner_mask
    return complete_re_pred`