SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix slow dpr test (#10059)

Browse Source 
* Correct cast to device

* Comment back the slow test
This commit is contained in:
Lysandre Debut
2021-02-08 10:43:25 +01:00
committed by GitHub
parent 12e44af5d3
commit d51302cca0
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_modeling_dpr.py
View File

@@ -265,6 +265,7 @@ class DPRModelIntegrationTest(unittest.TestCase):
def test_reader_inference(self): def test_reader_inference(self):
tokenizer = DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-single-nq-base") tokenizer = DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-single-nq-base")
model = DPRReader.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( encoded_inputs = tokenizer(
questions="What is love ?", questions="What is love ?",
@@ -273,6 +274,7 @@ class DPRModelIntegrationTest(unittest.TestCase):
padding=True, padding=True,
return_tensors="pt", return_tensors="pt",
) )
encoded_inputs.to(torch_device)
outputs = model(**encoded_inputs) outputs = model(**encoded_inputs)