test for model fetch

choderalab · Jul 14, 2023 · 4a8d970 · 4a8d970
1 parent 39fbd85
commit 4a8d970
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/espaloma/utils/tests/test_model_fetch.py b/espaloma/utils/tests/test_model_fetch.py
@@ -0,0 +1,15 @@
+import espaloma as esp
+import torch
+from espaloma.utils.model_fetch import get_model
+from openff.toolkit.topology import Molecule
+
+def test_get_model(tmp_path):
+    model_dir = tmp_path / "latest"
+    model_path = get_model(model_dir=model_dir, disable_progress_bar=True)
+
+    molecule = Molecule.from_smiles("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
+    molecule_graph = esp.Graph(molecule)
+
+    espaloma_model = torch.load(model_path)
+    espaloma_model.eval()
+    espaloma_model(molecule_graph.heterograph)