[CI] Fix partially red CI (#39448)

fix
2025-07-16 15:53:43 +02:00
parent 0a88751940
commit 26be7f717e
1 changed files with 1 additions and 2 deletions
--- a/tests/models/git/test_modeling_git.py
+++ b/tests/models/git/test_modeling_git.py
@@ -347,8 +347,7 @@ class GitModelTester:
            num_return_sequences=2,
        )

-        self.parent.assertEqual(generated_ids.shape[0], self.batch_size * 2)
-        self.parent.assertTrue(generated_ids.shape[1] < 20)
+        self.parent.assertEqual(generated_ids.shape, (self.batch_size * 2, 20))

    def _test_batched_generate_captioning(self, config, input_ids, input_mask, pixel_values):
        model = GitForCausalLM(config=config)