Fix bart slow test

2020-02-26 20:47:49 +00:00
parent 5bc99e7f33
commit f5516805c2
1 changed files with 1 additions and 1 deletions
--- a/tests/test_modeling_bart.py
+++ b/tests/test_modeling_bart.py
@@ -314,7 +314,7 @@ class BartModelIntegrationTest(unittest.TestCase):
            output = model.forward(**inputs_dict)[0]
        expected_shape = torch.Size((1, 11, 1024))
        self.assertEqual(output.shape, expected_shape)
-        expected_slice = torch.Tensor(
+        expected_slice = torch.tensor(
            [[0.7144, 0.8143, -1.2813], [0.7144, 0.8143, -1.2813], [-0.0467, 2.5911, -2.1845]], device=torch_device
        )
        self.assertTrue(torch.allclose(output[:, :3, :3], expected_slice, atol=TOLERANCE))