fix tearDown

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Jan 26, 2024 · 2079866 · 2079866
1 parent 228fec3
commit 2079866
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/source/tests/pt/test_jit.py b/source/tests/pt/test_jit.py
@@ -97,6 +97,9 @@ def setUp(self):
         self.config["training"]["numb_steps"] = 10
         self.config["training"]["save_freq"] = 10
 
+    def tearDown(self):
+        JITTest.tearDown(self)
+
 
 @unittest.skip("hybrid not supported at the moment")
 class TestEnergyModelHybrid(unittest.TestCase, JITTest):
@@ -111,6 +114,9 @@ def setUp(self):
         self.config["training"]["numb_steps"] = 10
         self.config["training"]["save_freq"] = 10
 
+    def tearDown(self):
+        JITTest.tearDown(self)
+
 
 @unittest.skip("hybrid not supported at the moment")
 class TestEnergyModelHybrid2(unittest.TestCase, JITTest):
@@ -126,6 +132,9 @@ def setUp(self):
         self.config["training"]["numb_steps"] = 10
         self.config["training"]["save_freq"] = 10
 
+    def tearDown(self):
+        JITTest.tearDown(self)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py
@@ -50,6 +50,9 @@ def setUp(self):
         self.config["training"]["numb_steps"] = 1
         self.config["training"]["save_freq"] = 1
 
+    def tearDown(self) -> None:
+        DPTrainTest.tearDown(self)
+
 
 class TestEnergyModelDPA1(unittest.TestCase, DPTrainTest):
     def setUp(self):
@@ -63,6 +66,9 @@ def setUp(self):
         self.config["training"]["numb_steps"] = 1
         self.config["training"]["save_freq"] = 1
 
+    def tearDown(self) -> None:
+        DPTrainTest.tearDown(self)
+
 
 class TestEnergyModelDPA2(unittest.TestCase, DPTrainTest):
     def setUp(self):
@@ -85,6 +91,9 @@ def setUp(self):
         self.config["training"]["numb_steps"] = 1
         self.config["training"]["save_freq"] = 1
 
+    def tearDown(self) -> None:
+        DPTrainTest.tearDown(self)
+
 
 @unittest.skip("hybrid not supported at the moment")
 class TestEnergyModelHybrid(unittest.TestCase, DPTrainTest):
@@ -99,6 +108,9 @@ def setUp(self):
         self.config["training"]["numb_steps"] = 1
         self.config["training"]["save_freq"] = 1
 
+    def tearDown(self) -> None:
+        DPTrainTest.tearDown(self)
+
 
 if __name__ == "__main__":
     unittest.main()