diff --git a/mddatasetbuilder/datasetbuilder.py b/mddatasetbuilder/datasetbuilder.py
index 101da46..254becf 100644
--- a/mddatasetbuilder/datasetbuilder.py
+++ b/mddatasetbuilder/datasetbuilder.py
@@ -203,8 +203,7 @@ def _writestepmatrix(self, item):
                 # atom ID starts from 1
                 distances = step_atoms.get_distances(
                     atoma-1, range(len(step_atoms)), mic=True)
-                cutoffatomid = np.where(distances < self.cutoff)
-                cutoffatoms = step_atoms[cutoffatomid]
+                cutoffatoms = step_atoms[distances < self.cutoff]
                 symbols = cutoffatoms.get_chemical_symbols()
                 results.append(
                     (np.array([step, atoma]),