From 2079866fc8a4dfe27911d073604bb2f4e9c325fd Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Fri, 26 Jan 2024 15:55:49 -0500
Subject: [PATCH] fix tearDown

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 source/tests/pt/test_jit.py      |  9 +++++++++
 source/tests/pt/test_training.py | 12 ++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/source/tests/pt/test_jit.py b/source/tests/pt/test_jit.py
index d9d88ddbc1..f13dade183 100644
--- 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
index 2174fe92ec..574ca8688e 100644
--- 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()