databricks/dolly-v2-12b
databricks/dolly-v2-12b
1
0
Overview Code Discussions Pull Requests Wiki Activity

Update instruct_pipeline.py

Browse Source 
Send attention_mask to device
This commit is contained in:
Matthew Hayes 2023-04-25 21:21:53 +00:00 committed by huggingface-web
parent a7077365ca
commit d0aa7ea43d
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
instruct_pipeline.py
View File

@ -131,7 +131,7 @@ class InstructionTextGenerationPipeline(Pipeline):
generated_sequence = self.model.generate( generated_sequence = self.model.generate(
input_ids=input_ids.to(self.model.device), 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, pad_token_id=self.tokenizer.pad_token_id,
**generate_kwargs, **generate_kwargs,
) )