@torch.no_grad()
def _predict(self, loader):
    self.model.eval()
    preds = {'arcs': [], 'rels': [], 'probs': ([] if self.args.prob else None)}
    for batch in progress_bar(loader):
        (words, texts, *feats) = batch
        word_mask = words.ne(self.args.pad_index)
        mask = (word_mask if (len(words.shape) < 3) else word_mask.any((- 1)))
        mask[:, 0] = 0
        lens = mask.sum(1).tolist()
        (s_arc, s_rel) = self.model(words, feats)
        (arc_preds, rel_preds) = self.model.decode(s_arc, s_rel, mask, self.args.tree, self.args.proj)
        preds['arcs'].extend(arc_preds[mask].split(lens))
        preds['rels'].extend(rel_preds[mask].split(lens))
        if self.args.prob:
            preds['probs'].extend([prob[1:(i + 1), :(i + 1)].cpu() for (i, prob) in zip(lens, s_arc.softmax((- 1)).unbind())])
    preds['arcs'] = [seq.tolist() for seq in preds['arcs']]
    preds['rels'] = [self.REL.vocab[seq.tolist()] for seq in preds['rels']]
    return preds
@torch.no_grad()
def _predict(self, loader):
    self.model.eval()
    preds = {'arcs': [], 'rels': [], 'probs': ([] if self.args.prob else None)}
    for batch in progress_bar(loader):
        (words, texts, *feats) = batch
        word_mask = words.ne(self.args.pad_index)
        mask = (word_mask if (len(words.shape) < 3) else word_mask.any((- 1)))
        mask[:, 0] = 0
        lens = mask.sum(1).tolist()
        (s_arc, s_rel) = self.model(words, feats)
        (arc_preds, rel_preds) = self.model.decode(s_arc, s_rel, mask, self.args.tree, self.args.proj)
        preds['arcs'].extend(arc_preds[mask].split(lens))
        preds['rels'].extend(rel_preds[mask].split(lens))
        if self.args.prob:
            preds['probs'].extend([prob[1:(i + 1), :(i + 1)].cpu() for (i, prob) in zip(lens, s_arc.softmax((- 1)).unbind())])
    preds['arcs'] = [seq.tolist() for seq in preds['arcs']]
    preds['rels'] = [self.REL.vocab[seq.tolist()] for seq in preds['rels']]
    return preds