Add force to unittest of JIT (#326)

273c9fd6 · Gao, Xiang · Farhad Ramezanghorbani · a1adceb0 · 273c9fd6
Commit 273c9fd6 authored Oct 09, 2019 by Gao, Xiang Committed by Farhad Ramezanghorbani Oct 09, 2019
Show whitespace changes
Inline Side-by-side

Showing with 7 additions and 0 deletions

tests/test_forces.py tests/test_forces.py +7 -0

No files found.
--- a/tests/test_forces.py
+++ b/tests/test_forces.py
@@ -67,5 +67,12 @@ class TestForce(unittest.TestCase):
            self.assertLess(max_diff, self.tolerance)


+class TestForceJIT(TestForce):
+
+    def setUp(self):
+        super().setUp()
+        self.model = torch.jit.script(self.model)
+
+
 if __name__ == '__main__':
    unittest.main()