fix test_generated_length_assisted_generation (#34935)

fix test_generated_length_assisted_generation
huggingface · Jan 29, 2025 · 42c8ccf · 42c8ccf
1 parent ec7afad
commit 42c8ccf
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py
@@ -3405,7 +3405,14 @@ def test_generated_length_assisted_generation(self):
             assistant_model=assistant,
             min_new_tokens=10,
         )
-        self.assertTrue((input_length + 10) <= out.shape[-1] <= 20)
+        self.assertTrue((input_length + 10) <= out.shape[-1])
+
+        out = model.generate(
+            input_ids,
+            assistant_model=assistant,
+            max_new_tokens=7,
+        )
+        self.assertTrue(out.shape[-1] <= (input_length + 7))
 
     def test_model_kwarg_assisted_decoding_decoder_only(self):
         # PT-only test: TF doesn't support assisted decoding yet.