fixed token_idx

huggingface · Dec 12, 2023 · e03d1af · e03d1af
1 parent 09286c3
commit e03d1af
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/habana/transformers/models/gptj/modeling_gptj.py b/optimum/habana/transformers/models/gptj/modeling_gptj.py
@@ -390,7 +390,7 @@ def gaudi_gptj_model_forward(
             def create_custom_forward(module):
                 def custom_forward(*inputs):
                     # None for past_key_value
-                    return module(*inputs, use_cache, output_attentions, token_idx, sin, cos)
+                    return module(*inputs, use_cache, output_attentions, None, sin, cos)
 
                 return custom_forward