inference update

Signed-off-by: Paarth Neekhara <[email protected]>
paarthneekhara · Nov 23, 2023 · 7e24168 · 7e24168
1 parent 271e2e3
commit 7e24168
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -1977,7 +1977,7 @@ def custom_autoregressive_inference(self, batch, prompt_len, pred_steps=500, sid
                     if _t % 10 == 0:
                         print("Decoding timestep", _t)
 
-                    (logits, _), _ = self.model(
+                    (logits, _), _, _ = self.model(
                         curr_tokens,
                         curr_position_ids,
                         curr_attention_mask,