From 864e9636ff2cd6ca1fad382f4e45bc8617b5cbec Mon Sep 17 00:00:00 2001 From: Fanli Lin Date: Thu, 24 Apr 2025 17:28:35 +0800 Subject: [PATCH] [tests] fix `test_nemotron_8b_generation_sdpa` (#37665) add max_new_tokens --- tests/models/nemotron/test_modeling_nemotron.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/nemotron/test_modeling_nemotron.py b/tests/models/nemotron/test_modeling_nemotron.py index d573537b6b..f7dcf27325 100644 --- a/tests/models/nemotron/test_modeling_nemotron.py +++ b/tests/models/nemotron/test_modeling_nemotron.py @@ -195,7 +195,7 @@ class NemotronIntegrationTest(unittest.TestCase): tokenizer = AutoTokenizer.from_pretrained(model_id) inputs = tokenizer(text, return_tensors="pt").to(torch_device) - output = model.generate(**inputs, do_sample=False) + output = model.generate(**inputs, do_sample=False, max_new_tokens=10) output_text = tokenizer.batch_decode(output, skip_special_tokens=True) self.assertEqual(EXPECTED_TEXT, output_text)