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

c58e6c12 · Sam Shleifer · GitHub · ddf9a3df · c58e6c12
Unverified Commit c58e6c12 authored Jun 06, 2020 by Sam Shleifer Committed by GitHub Jun 06, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

tests/test_modeling_marian.py tests/test_modeling_marian.py +2 -1

No files found.
--- a/tests/test_modeling_marian.py
+++ b/tests/test_modeling_marian.py
@@ -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])