Skip to content

Commit

Permalink
[bart] minor test fixes (#35965)
Browse files Browse the repository at this point in the history
fix tests
  • Loading branch information
gante authored Jan 30, 2025
1 parent 19f2ec8 commit 8bc4c89
Showing 2 changed files with 3 additions and 1 deletion.
2 changes: 1 addition & 1 deletion tests/models/bart/test_modeling_bart.py
Original file line number Diff line number Diff line change
@@ -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)

2 changes: 2 additions & 0 deletions tests/test_modeling_common.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 8bc4c89

Please sign in to comment.