BeamContainer(elem["score"].item(), elem) for elem in finalized[sent]
]
BCList.sort()
BCList.reverse()
finalized[sent] = torch.jit.annotate(
List[Dict[str, Tensor]], [x.elem for x in BCList]
)
After Change
for sent in range(len(finalized)):
scores = torch.tensor([float(elem["score"].item()) for elem in finalized[sent]])
_, sorted_scores_indices = torch.sort(scores, descending=True)
finalized[sent] = [finalized[sent][ssi] for ssi in sorted_scores_indices]
finalized[sent] = torch.jit.annotate(List[Dict[str, Tensor]], finalized[sent])
return finalized
def _prefix_tokens(