diff --git a/tests/models/patchtsmixer/test_modeling_patchtsmixer.py b/tests/models/patchtsmixer/test_modeling_patchtsmixer.py index 36a56fd1a1a298..70de9e516f23ac 100644 --- a/tests/models/patchtsmixer/test_modeling_patchtsmixer.py +++ b/tests/models/patchtsmixer/test_modeling_patchtsmixer.py @@ -469,18 +469,7 @@ def test_pretrain_head(self): ) self.assertEqual(output.shape, expected_shape) - expected_slice = torch.tensor( - [ - [[-0.9106]], - [[1.5326]], - [[-0.8245]], - [[0.7439]], - [[-0.7830]], - [[2.6256]], - [[-0.6485]], - ], - device=torch_device, - ) + expected_slice = torch.tensor([[[[-0.9106]],[[1.5326]],[[-0.8245]],[[0.7439]],[[-0.7830]],[[2.6256]],[[-0.6485]],]],device=torch_device) # fmt: skip self.assertTrue(torch.allclose(output[0, :7, :1, :1], expected_slice, atol=TOLERANCE)) def test_forecasting_head(self):