Fixed bug in notebook (NVIDIA#5382) (NVIDIA#5394)

Signed-off-by: Virginia Adams <vadams@nvidia.com> Signed-off-by: Virginia Adams <vadams@nvidia.com> Signed-off-by: Virginia Adams <vadams@nvidia.com> Co-authored-by: Virginia Adams <78445382+vadam5@users.noreply.github.com> Signed-off-by: andrusenkoau <andrusenkoau@gmail.com>
andrusenkoau · Jan 5, 2023 · 3e62d82 · 3e62d82
1 parent 9f1eb26
commit 3e62d82
Showing 1 changed file with 1 addition and 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
@@ -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: