[ci] fix 3 remaining slow GPU failures (#4584)

2020-05-25 19:20:50 -04:00
parent 365d452d4d
commit b86e42e0ac
4 changed files with 5 additions and 5 deletions
--- a/tests/test_modeling_bart.py
+++ b/tests/test_modeling_bart.py
@@ -240,7 +240,7 @@ class BartTranslationTests(unittest.TestCase):
        with torch.no_grad():
            logits, *other_stuff = model(**self.net_input)

-        expected_slice = torch.tensor([9.0078, 10.1113, 14.4787])
+        expected_slice = torch.tensor([9.0078, 10.1113, 14.4787], device=torch_device)
        result_slice = logits[0][0][:3]
        self.assertTrue(torch.allclose(expected_slice, result_slice, atol=TOLERANCE))