accept attention_mask in generate()

modeling_swa.py

@@ -174,7 +174,7 @@ class SWAForCausalLM(PreTrainedModel, GenerationMixin):
     def get_input_embeddings(s): return s.model.emb
     def set_input_embeddings(s, e): s.model.emb = e
 
-    def prepare_inputs_for_generation(s, input_ids, past_key_values=None, **kw):
+    def prepare_inputs_for_generation(s, input_ids, past_key_values=None, attention_mask=None, **kwargs):
         """generate(): keep one SWACache per sequence and feed only the tokens it has not seen. The cache is found in
         past_key_values, or (if this transformers version keeps its own cache object there) in the last cache this
         model built, reused only when its tokens are exactly the prefix of input_ids. Batch rows must be unpadded."""