Avoid loading/padding forces when not needed (#463)

* cleanup * move forces to additional_properties * fix typo
aiqm · May 1, 2020 · de9020d · de9020d
1 parent 4f30aae
commit de9020d
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 9 deletions.
diff --git a/examples/nnp_training_force.py b/examples/nnp_training_force.py
@@ -49,7 +49,11 @@
 
 batch_size = 2560
 
-training, validation = torchani.data.load(dspath).subtract_self_energies(energy_shifter).species_to_indices(species_order).shuffle().split(0.8, None)
+training, validation = torchani.data.load(
+    dspath,
+    additional_properties=('forces',)
+).subtract_self_energies(energy_shifter).species_to_indices(species_order).shuffle().split(0.8, None)
+
 training = training.collate(batch_size).cache()
 validation = validation.collate(batch_size).cache()
 

diff --git a/tests/test_jit_builtin_models.py b/tests/test_jit_builtin_models.py
@@ -6,13 +6,6 @@
 
 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):

diff --git a/torchani/data/__init__.py b/torchani/data/__init__.py
@@ -67,7 +67,7 @@
 verbose = True
 
 
-PROPERTIES = ('energies', 'forces')
+PROPERTIES = ('energies',)
 PADDING = {
     'species': -1,
     'coordinates': 0.0,