diff --git a/source/tests/universal/pt/model/test_model.py b/source/tests/universal/pt/model/test_model.py
index f17e12dc50..65bd981ccf 100644
--- a/source/tests/universal/pt/model/test_model.py
+++ b/source/tests/universal/pt/model/test_model.py
@@ -885,8 +885,7 @@ def setUpClass(cls):
             type_map=cls.expected_type_map,
         )
         cls.module = DPLinearModel(
-            dp_model,
-            pt_model,
+            [dp_model1,dp_model2],
             type_map=cls.expected_type_map,
         )
         # only test jit API once for different models