Pass `use_cache` in kwargs for GPTNeoX (#30538)

pass use_cache in kwargs

Pass `use_cache` in kwargs for GPTNeoX (#30538)
pass use_cache in kwargs
c712d05a · Raushan Turganbay · GitHub · a3aabc70 · c712d05a
Unverified Commit c712d05a authored Apr 30, 2024 by Raushan Turganbay Committed by GitHub Apr 30, 2024
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

src/transformers/models/gpt_neox/modeling_gpt_neox.py src/transformers/models/gpt_neox/modeling_gpt_neox.py +1 -0

No files found.
--- a/src/transformers/models/gpt_neox/modeling_gpt_neox.py
+++ b/src/transformers/models/gpt_neox/modeling_gpt_neox.py
@@ -1108,6 +1108,7 @@ class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel):
                "attention_mask": attention_mask,
                "past_key_values": past_key_values,
                "position_ids": position_ids,
+                "use_cache": kwargs.get("use_cache"),
            }
        )