fix for when uneven num of candidates
@@ -196,7 +196,11 @@ class TorchRankerAgent(TorchAgent):
cand_list = cands
elif cand_vecs.dim() == 3:
cand_list = cands[i]
+ if len(ordering) != len(cand_list):
+ true_ordering = [x for x in ordering if x < len(cand_list)]
+ ordering = true_ordering
cand_preds.append([cand_list[rank] for rank in ordering])
+
preds = [cand_preds[i][0] for i in range(batchsize)]
return Output(preds, cand_preds)
