Applied review comment

Signed-off-by: Taylor Yeonbok Lee <[email protected]>

Author: taylor-yb-lee

tests/unittest/llmapi/test_llm_args.py

@@ -138,28 +138,17 @@ def test_llm_args_with_pydantic_options(self):
         assert llm_args.max_num_tokens == 256
         assert llm_args.max_seq_len == 128
 
-    def test_llm_args_with_model_kwargs_trt(self):
+    @pytest.mark.parametrize("llm_args_cls", [TrtLlmArgs, TorchLlmArgs])
+    def test_llm_args_with_model_kwargs(self, llm_args_cls):
         yaml_content = """
 model_kwargs:
     num_hidden_layers: 2
     """
         dict_content = self._yaml_to_dict(yaml_content)
-        llm_args = TrtLlmArgs(model=llama_model_path)
-        llm_args_dict = update_llm_args_with_extra_dict(llm_args.model_dump(),
-                                                        dict_content)
-        llm_args = TrtLlmArgs(**llm_args_dict)
-        assert llm_args.model_kwargs['num_hidden_layers'] == 2
-
-    def test_llm_args_with_model_kwargs_pt(self):
-        yaml_content = """
-model_kwargs:
-    num_hidden_layers: 2
-    """
-        dict_content = self._yaml_to_dict(yaml_content)
-        llm_args = TorchLlmArgs(model=llama_model_path)
+        llm_args = llm_args_cls(model=llama_model_path)
         llm_args_dict = update_llm_args_with_extra_dict(llm_args.model_dump(),
                                                         dict_content)
-        llm_args = TorchLlmArgs(**llm_args_dict)
+        llm_args = llm_args_cls(**llm_args_dict)
         assert llm_args.model_kwargs['num_hidden_layers'] == 2