fix: UTs

anyangml · Oct 8, 2024 · f2753e7 · f2753e7
1 parent 8b1cb8c
commit f2753e7
Showing 1 changed file with 1 addition and 2 deletions.
diff --git 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