Patch T5 device test (#12742)

2021-07-15 17:40:17 +02:00
parent 370be9cc38
commit f42d9dcc0e
1 changed files with 1 additions and 1 deletions
--- a/tests/test_modeling_t5.py
+++ b/tests/test_modeling_t5.py
@@ -802,7 +802,7 @@ class T5ModelIntegrationTests(unittest.TestCase):
        model.config.do_sample = False
        tokenizer = T5Tokenizer.from_pretrained("t5-small")

-        input_ids = tokenizer("summarize: Hello there", return_tensors="pt").input_ids
+        input_ids = tokenizer("summarize: Hello there", return_tensors="pt").input_ids.to(torch_device)

        sequences = model.generate(input_ids)