Pass use_cache in kwargs for GPTNeoX (#30538)

pass use_cache in kwargs
huggingface · May 14, 2024 · e52ffbe · e52ffbe
1 parent 2782292
commit e52ffbe
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/transformers/models/gpt_neox/modeling_gpt_neox.py b/src/transformers/models/gpt_neox/modeling_gpt_neox.py
@@ -1108,6 +1108,7 @@ def prepare_inputs_for_generation(
                 "attention_mask": attention_mask,
                 "past_key_values": past_key_values,
                 "position_ids": position_ids,
+                "use_cache": kwargs.get("use_cache"),
             }
         )