From f5516805c2d0ea39797d46d9433dab43f769bea1 Mon Sep 17 00:00:00 2001 From: Julien Chaumond Date: Wed, 26 Feb 2020 20:47:49 +0000 Subject: [PATCH] Fix bart slow test --- tests/test_modeling_bart.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_modeling_bart.py b/tests/test_modeling_bart.py index ef8932618a..29b459fd8d 100644 --- 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))