riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-31 02:02:21 +06:00
Code Issues Actions 23 Packages Projects Releases Wiki Activity

[marian tests ] pass device to pipeline (#4815)

Browse Source
This commit is contained in:
Sam Shleifer 2020-06-06 00:52:17 -04:00 committed by GitHub
parent ddf9a3dfc7
commit c58e6c129a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/test_modeling_marian.py
View File

@ -233,7 +233,8 @@ class TestMarian_en_ROMANCE(MarianIntegrationTest):
self.tokenizer.prepare_translation_batch([""])
def test_pipeline(self):
pipeline = TranslationPipeline(self.model, self.tokenizer, framework="pt")
device = 0 if torch_device == "cuda" else -1
pipeline = TranslationPipeline(self.model, self.tokenizer, framework="pt", device=device)
output = pipeline(self.src_text)
self.assertEqual(self.expected_text, [x["translation_text"] for x in output])