SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 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
1 changed files with 2 additions and 1 deletions
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([""]) self.tokenizer.prepare_translation_batch([""])
def test_pipeline(self): 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) output = pipeline(self.src_text)
self.assertEqual(self.expected_text, [x["translation_text"] for x in output]) self.assertEqual(self.expected_text, [x["translation_text"] for x in output])