From 1b6386be89b1f57a3ff69523aeda55604db369e2 Mon Sep 17 00:00:00 2001 From: Yi Dong Date: Mon, 23 May 2022 20:37:40 +0000 Subject: [PATCH] fix json serialize Signed-off-by: Yi Dong --- nemo/collections/nlp/modules/common/text_generation_server.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/nemo/collections/nlp/modules/common/text_generation_server.py b/nemo/collections/nlp/modules/common/text_generation_server.py index 40c9dc385e5e..3939f82f3e0d 100644 --- a/nemo/collections/nlp/modules/common/text_generation_server.py +++ b/nemo/collections/nlp/modules/common/text_generation_server.py @@ -158,6 +158,9 @@ def put(self): repetition_penalty, min_tokens_to_generate, ) + for k in output: + if isinstance(output[k], torch.Tensor): + output[k] = output[k].tolist() if not all_probs: del output['full_logprob'] return jsonify(output)