diff --git a/parlai/core/torch_classifier_agent.py b/parlai/core/torch_classifier_agent.py index 786e79607de..6ac2471147a 100644 --- a/parlai/core/torch_classifier_agent.py +++ b/parlai/core/torch_classifier_agent.py @@ -703,7 +703,7 @@ def eval_step(self, batch): self._update_confusion_matrix(batch, preds) if self.opt.get('print_scores', False): - return Output(preds, class_list=[self.class_list], probs=probs.cpu()) + return Output(preds, class_list=[self.class_list], probs=probs.tolist()) if self.opt.get('return_cand_scores', False): sorted_scores, ranks = probs.sort(1, descending=True) sorted_scores = sorted_scores.cpu()