xxxxx

Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Oct 8, 2024 · bf6db03 · bf6db03
1 parent 013883c
commit bf6db03
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/source/tests/universal/dpmodel/backend.py b/source/tests/universal/dpmodel/backend.py
@@ -41,7 +41,9 @@ def _get_deserialized_module(cls):
 
     @property
     def deserialized_module(self):
-        return self._get_deserialized_module()
+        if hasattr(self.__class__, "module"):
+            return self._get_deserialized_module()
+        return self.module.deserialize(self.module.serialize())
 
     @property
     def modules_to_test(self):

diff --git a/source/tests/universal/pt/backend.py b/source/tests/universal/pt/backend.py
@@ -30,7 +30,10 @@ def _get_script_module(cls):
 
     @property
     def script_module(self):
-        return self._get_script_module()
+        if hasattr(self.__class__, "module"):
+            return self._get_script_module()
+        with torch.jit.optimized_execution(False):
+            return torch.jit.script(self.module)
 
     @classmethod
     @cache
@@ -39,7 +42,9 @@ def _get_deserialized_module(cls):
 
     @property
     def deserialized_module(self):
-        return self._get_deserialized_module()
+        if hasattr(self.__class__, "module"):
+            return self._get_deserialized_module()
+        return self.module.deserialize(self.module.serialize())
 
     @property
     def modules_to_test(self):