Unverified Commit d51302cc authored by Lysandre Debut's avatar Lysandre Debut Committed by GitHub
Browse files

Fix slow dpr test (#10059)

* Correct cast to device

* Comment back the slow test
parent 12e44af5
......@@ -265,6 +265,7 @@ class DPRModelIntegrationTest(unittest.TestCase):
def test_reader_inference(self):
tokenizer = DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-single-nq-base")
model = DPRReader.from_pretrained("facebook/dpr-reader-single-nq-base")
model.to(torch_device)
encoded_inputs = tokenizer(
questions="What is love ?",
......@@ -273,6 +274,7 @@ class DPRModelIntegrationTest(unittest.TestCase):
padding=True,
return_tensors="pt",
)
encoded_inputs.to(torch_device)
outputs = model(**encoded_inputs)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment