Add force to unittest of JIT (#326)

aiqm · Oct 9, 2019 · 273c9fd · 273c9fd
1 parent a1adceb
commit 273c9fd
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/tests/test_forces.py b/tests/test_forces.py
@@ -67,5 +67,12 @@ def testPadding(self):
             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()