From 8f3ca3802b569953fc5d756aeef465558bbef5fe Mon Sep 17 00:00:00 2001 From: Virginia Adams Date: Wed, 9 Nov 2022 23:54:50 +0000 Subject: [PATCH] Fixed bug in notebook Signed-off-by: Virginia Adams --- .../language_modeling/megatron_gpt_prompt_learning_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_prompt_learning_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_prompt_learning_model.py index add7c898c80c..387ff52bb078 100644 --- a/nemo/collections/nlp/models/language_modeling/megatron_gpt_prompt_learning_model.py +++ b/nemo/collections/nlp/models/language_modeling/megatron_gpt_prompt_learning_model.py @@ -986,7 +986,7 @@ def dummy(): # Call same generate code as in MegatronGPT return megatron_gpt_generate( - self.cuda(), processed_inputs, self.tokenizer, length_params, sampling_params, task_ids + self.cuda(), processed_inputs, self.tokenizer, length_params, sampling_params, task_ids=task_ids ) def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: Optional[int] = None) -> Any: