From d51302cca0eb46f32a067aedc91ffee3edc29acb Mon Sep 17 00:00:00 2001
From: Lysandre Debut <lysandre@huggingface.co>
Date: Mon, 8 Feb 2021 10:43:25 +0100
Subject: [PATCH] Fix slow dpr test (#10059)

* Correct cast to device

* Comment back the slow test
---
 tests/test_modeling_dpr.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/test_modeling_dpr.py b/tests/test_modeling_dpr.py
index f31563670a..05c9844b4b 100644
--- a/tests/test_modeling_dpr.py
+++ b/tests/test_modeling_dpr.py
@@ -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)