From 5e6880724a61a61d4f2d982471a05c5e72ca8562 Mon Sep 17 00:00:00 2001 From: Jason Weston Date: Wed, 5 Jan 2022 18:16:29 -0500 Subject: [PATCH] TCA: probability bug --- parlai/core/torch_classifier_agent.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/parlai/core/torch_classifier_agent.py b/parlai/core/torch_classifier_agent.py index 84e4b734a4a..baa4aa20c74 100644 --- a/parlai/core/torch_classifier_agent.py +++ b/parlai/core/torch_classifier_agent.py @@ -705,7 +705,9 @@ def eval_step(self, batch): for i in range(0, ranks.size(0)): ordered_list = [self.class_list[i] for i in ranks[i]] text_cands.append(ordered_list) - return Output(preds, text_candidates=text_cands, sorted_scores=probs.cpu()) + return Output( + preds, text_candidates=text_cands, sorted_scores=sorted_scores + ) else: return Output(preds)