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

This commit is contained in:
Alberto Villa 2021-06-01 20:58:49 +02:00 committed by GitHub
parent 47a98fc4cb
commit 1c3ab3e5d6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -90,7 +90,7 @@ Usage Example
>>> device = 'cuda' if torch.cuda.is_available() else 'cpu'
>>> tokenizer = PegasusTokenizer.from_pretrained(model_name)
>>> 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)
>>> 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."