SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Typo in usage example, changed to device instead of torch_device (#11979)

Browse Source
This commit is contained in:
Alberto Villa
2021-06-01 20:58:49 +02:00
committed by GitHub
parent 47a98fc4cb
commit 1c3ab3e5d6
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/source/model_doc/pegasus.rst
View File

@@ -90,7 +90,7 @@ Usage Example
>>> device = 'cuda' if torch.cuda.is_available() else 'cpu' >>> device = 'cuda' if torch.cuda.is_available() else 'cpu'
>>> tokenizer = PegasusTokenizer.from_pretrained(model_name) >>> tokenizer = PegasusTokenizer.from_pretrained(model_name)
>>> model = PegasusForConditionalGeneration.from_pretrained(model_name).to(device) >>> model = PegasusForConditionalGeneration.from_pretrained(model_name).to(device)
>>> batch = tokenizer(src_text, truncation=True, padding='longest', return_tensors="pt").to(torch_device) >>> batch = tokenizer(src_text, truncation=True, padding='longest', return_tensors="pt").to(device)
>>> translated = model.generate(**batch) >>> translated = model.generate(**batch)
>>> tgt_text = tokenizer.batch_decode(translated, skip_special_tokens=True) >>> tgt_text = tokenizer.batch_decode(translated, skip_special_tokens=True)
>>> assert tgt_text[0] == "California's largest electricity provider has turned off power to hundreds of thousands of customers." >>> assert tgt_text[0] == "California's largest electricity provider has turned off power to hundreds of thousands of customers."