Send attention_mask to device (#9)

- Send attention_mask to device (76fb033a3c0c4b1d6764fcf69699ba60f6ad942e)


Co-authored-by: Michael Verrilli <[email protected]>

instruct_pipeline.py

@@ -131,7 +131,7 @@ class InstructionTextGenerationPipeline(Pipeline):
 
         generated_sequence = self.model.generate(
             input_ids=input_ids.to(self.model.device),
-            attention_mask=attention_mask,
+            attention_mask=attention_mask.to(self.model.device) if attention_mask is not None else None,
             pad_token_id=self.tokenizer.pad_token_id,
             **generate_kwargs,
         )