From 7e24168d207f5fa3c77c4b7ed8d9839ddbf17785 Mon Sep 17 00:00:00 2001
From: Paarth Neekhara <paarth.n@gmail.com>
Date: Wed, 22 Nov 2023 23:08:52 -0500
Subject: [PATCH] inference update

Signed-off-by: Paarth Neekhara <paarth.n@gmail.com>
---
 .../nlp/models/language_modeling/megatron_gpt_model.py          | 2 +-
 1 file changed, 1 insertion(+), 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
index cea1b9fa8a56..84f59a0ce883 100644
--- 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,