Masking: remove flakiness from test (huggingface#31939)

amyeroberts · Jul 19, 2024 · bf99bda · bf99bda
1 parent 45d0d13
commit bf99bda
Show file tree

Hide file tree

Showing 2 changed files with 0 additions and 6 deletions.
diff --git a/tests/models/whisper/test_modeling_whisper.py b/tests/models/whisper/test_modeling_whisper.py
@@ -1571,9 +1571,6 @@ def test_custom_4d_attention_mask(self):
         out_last_tokens = logits[:, -1, :]  # last tokens in each batch line
         out_shared_prefix_last_tokens = logits_shared_prefix[0, -3:, :]  # last three tokens
 
-        # comparing greedily-chosen tokens:
-        assert torch.equal(out_last_tokens.max(axis=1).indices, out_shared_prefix_last_tokens.max(axis=1).indices)
-
         # comparing softmax-normalized logits:
         normalized_0 = torch.nn.functional.softmax(out_last_tokens)
         normalized_1 = torch.nn.functional.softmax(out_shared_prefix_last_tokens)

diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
@@ -4486,9 +4486,6 @@ def test_custom_4d_attention_mask(self):
             out_last_tokens = logits[:, -1, :]  # last tokens in each batch line
             out_shared_prefix_last_tokens = logits_shared_prefix[0, -3:, :]  # last three tokens
 
-            # comparing greedily-chosen tokens:
-            assert torch.equal(out_last_tokens.max(axis=1).indices, out_shared_prefix_last_tokens.max(axis=1).indices)
-
             # comparing softmax-normalized logits:
             normalized_0 = F.softmax(out_last_tokens)
             normalized_1 = F.softmax(out_shared_prefix_last_tokens)