Fix bug in discarding outlier energy conformers (#334)

aiqm · Oct 16, 2019 · b8c2aa9 · b8c2aa9
1 parent 10c723b
commit b8c2aa9
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/torchani/data/__init__.py b/torchani/data/__init__.py
@@ -310,13 +310,12 @@ def load_ani_dataset(path, species_tensor_converter, batch_size, shuffle=True,
         atomic_properties_, properties_ = t(atomic_properties_, properties_)
 
     if rm_outlier:
-        # This is how NeuroChem discard the outliers
         transformed_energies = properties_['energies']
         num_atoms = (atomic_properties_['species'] >= 0).to(transformed_energies.dtype).sum(dim=1)
         scaled_diff = transformed_energies / num_atoms.sqrt()
 
         mean = scaled_diff[torch.abs(scaled_diff) < 15.0].mean()
-        std = torch.abs(scaled_diff[torch.abs(scaled_diff) < 15.0]).std()
+        std = scaled_diff[torch.abs(scaled_diff) < 15.0].std()
 
         # -15 * std + mean < scaled_diff < +11 * std + mean
         tol = 13.0 * std + mean