fix the testsuit

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: yd-oom

tests/test_draft_modeling_tp.py

@@ -234,7 +234,7 @@ def run_attention_tp_test(rank, world_size, temp_dir_name):
         else:
             sharded_param = param
         sharded_state_dict[name] = sharded_param
-    attn_tp2.load_state_dict(sharded_state_dict, strict=False)
+    attn_tp2.load_state_dict(sharded_state_dict)
     attn_tp2.eval()
 
     input_tensor = torch.load(os.path.join(temp_dir_name, "attn_input.pth")).cuda(rank)