use proper config

Signed-off-by: Gerald Shen <[email protected]>
NVIDIA · Jul 12, 2023 · a05ac06 · a05ac06
1 parent 5416242
commit a05ac06
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
@@ -755,7 +755,7 @@ def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: Optional[int]
                 "all_probs": inference_config["all_probs"],
                 "compute_logprob": inference_config["compute_logprob"],
                 "compute_attention_mask": inference_config.get("compute_attention_mask", True),
-                "end_strings": self.cfg.inference.get('end_strings', ["<|endoftext|>"]),
+                "end_strings": inference_config.get('end_strings', ["<|endoftext|>"]),
             }
 
             task_ids, processed_inputs = batch