From 8bc4c89ee904598d402851f06d681bc1e9cc5846 Mon Sep 17 00:00:00 2001 From: Joao Gante Date: Thu, 30 Jan 2025 10:00:11 +0000 Subject: [PATCH] [bart] minor test fixes (#35965) fix tests --- tests/models/bart/test_modeling_bart.py | 2 +- tests/test_modeling_common.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/models/bart/test_modeling_bart.py b/tests/models/bart/test_modeling_bart.py index a795bfcabfb4..6c9db62daec2 100644 --- a/tests/models/bart/test_modeling_bart.py +++ b/tests/models/bart/test_modeling_bart.py @@ -885,7 +885,7 @@ def test_inference_no_head(self): expected_shape = torch.Size((1, 11, 1024)) self.assertEqual(output.shape, expected_shape) 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 + [[[0.7144, 0.8143, -1.2813], [0.7144, 0.8143, -1.2813], [-0.0467, 2.5911, -2.1845]]], device=torch_device ) torch.testing.assert_close(output[:, :3, :3], expected_slice, rtol=1e-3, atol=1e-3) diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 1d52e8281449..3f00b4e15d4b 100755 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -4532,6 +4532,8 @@ def test_flash_attention_2_padding_matches_padding_free_with_position_ids(self): config.max_position_embeddings = max_new_tokens + dummy_input.shape[1] + 1 model = model_class(config) + if "position_ids" not in inspect.signature(model.forward).parameters: + self.skipTest("Model does not support position_ids") with tempfile.TemporaryDirectory() as tmpdirname: model.save_pretrained(tmpdirname)