Text Generation
Transformers
PyTorch
mosaic_gpt
custom_code
anas-awadalla commited on
Commit
a453954
1 Parent(s): a2d0bd8

Update mosaic_gpt.py

Browse files
Files changed (1) hide show
  1. mosaic_gpt.py +2 -1
mosaic_gpt.py CHANGED
@@ -392,6 +392,7 @@ class MosaicGPT(PreTrainedModel):
392
 
393
  def prepare_inputs_for_generation(self,
394
  input_ids,
 
395
  past_key_values=None,
396
  inputs_embeds=None,
397
  **kwargs):
@@ -399,7 +400,7 @@ class MosaicGPT(PreTrainedModel):
399
  raise NotImplementedError(
400
  'inputs_embeds is not implemented for MosaicGPT yet')
401
 
402
- attention_mask = kwargs['attention_mask'].bool()
403
  if attention_mask[:, -1].sum() != attention_mask.shape[0]:
404
  raise NotImplementedError(
405
  'MosaicGPT does not support generation with right padding.')
 
392
 
393
  def prepare_inputs_for_generation(self,
394
  input_ids,
395
+ attention_mask=None,
396
  past_key_values=None,
397
  inputs_embeds=None,
398
  **kwargs):
 
400
  raise NotImplementedError(
401
  'inputs_embeds is not implemented for MosaicGPT yet')
402
 
403
+ attention_mask = attention_mask.bool()
404
  if attention_mask[:, -1].sum() != attention_mask.shape[0]:
405
  raise NotImplementedError(
406
  'MosaicGPT does not support generation with right padding.')