Update mosaic_gpt.py

mosaic_gpt.py

@@ -392,6 +392,7 @@ class MosaicGPT(PreTrainedModel):
 
     def prepare_inputs_for_generation(self,
                                       input_ids,
+                                      attention_mask=None,
                                       past_key_values=None,
                                       inputs_embeds=None,
                                       **kwargs):
@@ -399,7 +400,7 @@ class MosaicGPT(PreTrainedModel):
             raise NotImplementedError(
                 'inputs_embeds is not implemented for MosaicGPT yet')
 
-        attention_mask = kwargs['attention_mask'].bool()
+        attention_mask = attention_mask.bool()
         if attention_mask[:, -1].sum() != attention_mask.shape[0]:
             raise NotImplementedError(
                 'MosaicGPT does not support generation with right padding.')