SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[tests] fix test_nemotron_8b_generation_sdpa (#37665)

Browse Source 
add max_new_tokens
This commit is contained in:
Fanli Lin
2025-04-24 17:28:35 +08:00
committed by GitHub
parent 9b3bf4a206
commit 864e9636ff
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/models/nemotron/test_modeling_nemotron.py
View File

@@ -195,7 +195,7 @@ class NemotronIntegrationTest(unittest.TestCase):
tokenizer = AutoTokenizer.from_pretrained(model_id) tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer(text, return_tensors="pt").to(torch_device) 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) output_text = tokenizer.batch_decode(output, skip_special_tokens=True)
self.assertEqual(EXPECTED_TEXT, output_text) self.assertEqual(EXPECTED_TEXT, output_text)