From d1ce6f07f6fc304da1bd751e803678d2531b17c9 Mon Sep 17 00:00:00 2001
From: Xiang Gao <qasdfgtyuiop@gmail.com>
Date: Wed, 9 Oct 2019 13:50:16 -0700
Subject: [PATCH 1/2] Add unit test for builtin models

---
 .github/workflows/unittest.yml   |  6 +++-
 tests/test_jit_builtin_models.py | 49 ++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_jit_builtin_models.py

diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml
index ae6349e4c..8e70f9564 100644
--- a/.github/workflows/unittest.yml
+++ b/.github/workflows/unittest.yml
@@ -10,7 +10,11 @@ jobs:
       fail-fast: false
       matrix:
         python-version: [3.6, 3.7]
-        test-filenames: [test_aev.py, test_aev_benzene_md.py, test_aev_nist.py, test_aev_tripeptide_md.py, test_data_new.py, test_ignite.py, test_utils.py, test_ase.py, test_energies.py, test_neurochem.py, test_vibrational.py, test_ensemble.py, test_padding.py, test_data.py, test_forces.py, test_structure_optim.py]
+        test-filenames: [
+          test_aev.py, test_aev_benzene_md.py, test_aev_nist.py, test_aev_tripeptide_md.py,
+          test_data_new.py, test_ignite.py, test_utils.py, test_ase.py, test_energies.py,
+          test_neurochem.py, test_vibrational.py, test_ensemble.py, test_padding.py,
+          test_data.py, test_forces.py, test_structure_optim.py, test_jit_builtin_models.py]
 
     steps:
     - uses: actions/checkout@v1
diff --git a/tests/test_jit_builtin_models.py b/tests/test_jit_builtin_models.py
new file mode 100644
index 000000000..810900b2d
--- /dev/null
+++ b/tests/test_jit_builtin_models.py
@@ -0,0 +1,49 @@
+import torch
+import torchani
+import unittest
+import os
+import pickle
+
+
+path = os.path.dirname(os.path.realpath(__file__))
+dspath = os.path.join(path, '../dataset/ani-1x/sample.h5')
+batch_size = 256
+chunk_threshold = 5
+other_properties = {'properties': ['energies'],
+                    'padding_values': [None],
+                    'padded_shapes': [(batch_size, )],
+                    'dtypes': [torch.float64],
+                    }
+
+
+class TestBuiltinModelsJIT(unittest.TestCase):
+
+    def setUp(self):
+        self.ds = torchani.data.CachedDataset(dspath, batch_size=batch_size, device='cpu',
+                                              chunk_threshold=chunk_threshold,
+                                              other_properties=other_properties,
+                                              subtract_self_energies=True)
+        self.ani1ccx = torchani.models.ANI1ccx()
+
+    def _test_model(self, model):
+        chunk = self.ds[0][0][0]
+        _, e = model(chunk)
+        _, e2 = torch.jit.script(model)(chunk)
+        self.assertTrue(torch.allclose(e, e2))
+
+    def _test_ensemble(self, ensemble):
+        self._test_model(ensemble)
+        for m in ensemble:
+            self._test_model(m)
+
+    def testANI1x(self):
+        ani1x = torchani.models.ANI1x()
+        self._test_ensemble(ani1x)
+
+    def testANI1ccx(self):
+        ani1ccx = torchani.models.ANI1ccx()
+        self._test_ensemble(ani1ccx)
+
+
+if __name__ == '__main__':
+    unittest.main()

From a9499d3132528a89b52539eb33dcb9f7598a355f Mon Sep 17 00:00:00 2001
From: Xiang Gao <qasdfgtyuiop@gmail.com>
Date: Wed, 9 Oct 2019 15:22:35 -0700
Subject: [PATCH 2/2] flake8

---
 tests/test_jit_builtin_models.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/test_jit_builtin_models.py b/tests/test_jit_builtin_models.py
index 810900b2d..1d2814f19 100644
--- a/tests/test_jit_builtin_models.py
+++ b/tests/test_jit_builtin_models.py
@@ -2,7 +2,6 @@
 import torchani
 import unittest
 import os
-import pickle
 
 
 path = os.path.dirname(os.path.realpath(__file__))