Merge branch 'main' into coulomb

torchmd · Feb 20, 2024 · b619e0b · b619e0b
2 parents 1a73cb7 + 5a4fca7
commit b619e0b
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/examples/TensorNet-ANI2X.yaml b/examples/TensorNet-ANI2X.yaml
@@ -39,6 +39,7 @@ num_rbf: 32
 num_workers: 4
 output_model: Scalar
 precision: 32
+prior_model: null
 remove_ref_energy: true
 rbf_type: expnorm
 redirect: false

diff --git a/examples/openmm-integration.py b/examples/openmm-integration.py
@@ -51,8 +51,8 @@ def forward(self, positions):
 
 system = System()
 # Create an OpenMM system and add the TorchForce
-for i in range(pdb.topology.getNumAtoms()):
-    system.addParticle(1.0)
+for atom in pdb.topology.atoms():
+    system.addParticle(atom.element.mass)
 system.addForce(torch_force)
 integrator = LangevinMiddleIntegrator(298.15*kelvin, 1/picosecond, 2*femtosecond)
 platform = Platform.getPlatformByName('CPU')

diff --git a/tests/test_examples.py b/tests/test_examples.py
@@ -19,7 +19,7 @@ def test_example_yamls(fname):
         args = yaml.load(f, Loader=yaml.FullLoader)
 
     prior = None
-    if "prior_model" in args and args["prior_model"] is not None:
+    if args["prior_model"] is not None:
         dataset = DummyDataset(has_atomref=True)
         prior = getattr(priors, args["prior_model"])(dataset=dataset)