add 1 decode to test [skip ci]

Jackmin801 · Feb 7, 2025 · a061616 · a061616
1 parent 496b572
commit a061616
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/srt/test_srt_engine.py b/test/srt/test_srt_engine.py
@@ -191,7 +191,7 @@ def test_8_engine_return_hidden_states(self):
         tokenizer = AutoTokenizer.from_pretrained(model_path)
         input_ids = tokenizer(prompts).input_ids
 
-        sampling_params = {"temperature": 0, "max_new_tokens": 1}
+        sampling_params = {"temperature": 0, "max_new_tokens": 2}
 
         engine = sgl.Engine(
             model_path=model_path,
@@ -203,7 +203,7 @@ def test_8_engine_return_hidden_states(self):
         engine.shutdown()
 
         for output in outputs:
-            self.assertEqual(len(output["meta_info"]["hidden_states"]), 1)
+            self.assertEqual(len(output["meta_info"]["hidden_states"]), 2)
             for hidden_state in output["meta_info"]["hidden_states"]:
                 self.assertIsInstance(hidden_state, torch.Tensor)
         # Checks that splicing of the batch was done correctly