fix

Author: co63oc

tests/operators/test_eagle_get_hidden_states.py

@@ -26,15 +26,15 @@ def test_eagle_get_hidden_states(self):
         paddle.seed(2023)
         bs = np.random.randint(1, 8 + 1, dtype=np.int32)
         input_token_num = np.random.randint(2 * 1024, 4 * 1024 + 1, dtype=np.int32)
-        dim_embed = np.random.randint(1, 4 * 1024 + 1, dtype=np.int32)
+        dim_embed = np.array(1024, dtype=np.int32)
         actual_draft_token_num = np.random.randint(2, 6, dtype=np.int32)
 
         seq_lens_this_time = np.random.randint(0, 2, bs, dtype=np.int32)
         seq_lens_encoder = np.random.randint(0, input_token_num // bs + 1, bs, dtype=np.int32)
         accept_nums = np.random.randint(0, actual_draft_token_num + 1, bs, dtype=np.int32)
         base_model_seq_lens_this_time = np.random.randint(0, input_token_num // bs + 1, bs, dtype=np.int32)
         base_model_seq_lens_encoder = np.random.randint(0, 2, bs, dtype=np.int32)
-        # don't care
+
         seq_lens_decoder = np.random.randint(0, input_token_num // bs + 1, bs, dtype=np.int32)
         stop_flags = np.random.randint(0, 2, bs, dtype=np.int32)
 
@@ -43,13 +43,12 @@ def test_eagle_get_hidden_states(self):
         accept_nums_tensor = paddle.to_tensor(accept_nums, dtype=paddle.int32)
         base_model_seq_lens_this_time_tensor = paddle.to_tensor(base_model_seq_lens_this_time, dtype=paddle.int32)
         base_model_seq_lens_encoder_tensor = paddle.to_tensor(base_model_seq_lens_encoder, dtype=paddle.int32)
-        # don't care
+
         seq_lens_decoder_tensor = paddle.to_tensor(seq_lens_decoder, dtype=paddle.int32)
         stop_flags_tensor = paddle.to_tensor(stop_flags, dtype=paddle.int32)
 
-        # fp32 test
         input = np.random.randint(0, 10, (input_token_num, dim_embed), dtype=np.int32)
-        input_tensor = paddle.to_tensor(input, dtype=paddle.float32)
+        input_tensor = paddle.to_tensor(input, dtype=paddle.float16)
         gpu_out = eagle_get_hidden_states(
             input_tensor,
             seq_lens_this_time_tensor,
@@ -61,7 +60,8 @@ def test_eagle_get_hidden_states(self):
             base_model_seq_lens_encoder_tensor,
             actual_draft_token_num,
         )
-        print(gpu_out.numpy())
+        out_ref = np.array([6, 4, 3, 3], dtype=np.int32)
+        np.testing.assert_allclose(gpu_out.numpy()[0][0:4], out_ref)
 
 
 if __name__ == "__main__":