Update mosaic_gpt.py (#2)

- Update mosaic_gpt.py (150348683b5cfde8abc27f6a23b65a922e099fca)


Co-authored-by: Irena <[email protected]>

mosaic_gpt.py

@@ -247,7 +247,7 @@ class MosaicGPT(PreTrainedModel):
             use_cache: Optional[bool] = None):
         return_dict = return_dict if return_dict is not None else self.config.return_dict
         use_cache = use_cache if use_cache is not None else self.config.use_cache
-        attention_mask = attention_mask.bool()
+        attention_mask = attention_mask.bool() if attention_mask is not None else None
         
         # These args are passed in by keyword in huggingface's generate function
         # https://github.com/huggingface/transformers/blob/68287689f2f0d8b7063c400230b3766987abf18d/src/transformers/generation/utils.py#L2201-L2206