diff --git a/alframework/ml_interfaces/hippynn_interface.py b/alframework/ml_interfaces/hippynn_interface.py
index 6117180..cedf0e6 100644
--- a/alframework/ml_interfaces/hippynn_interface.py
+++ b/alframework/ml_interfaces/hippynn_interface.py
@@ -374,12 +374,7 @@ def train_HIPNN_model(model_dir,
             arrays = database.arr_dict
 
             arrays[species_key] = arrays[species_key].to(torch.int64)   # or .long()
-
-            min_dists = hippynn.pretraining.calculate_min_dists(database.arr_dict,species_key,coordinates_key,6.0,cell_name=cell_key,batch_size=256)
-            min_dists_flag = min_dists>network_params["dist_soft_min"]
-            for curK in database.arr_dict.keys():
-                database.arr_dict[curK] = database.arr_dict[curK][min_dists_flag]
-                
+ 
             n_atoms = (
               arrays[species_key]          # tensor
               .to(torch.bool)              # cast to bool
@@ -387,7 +382,7 @@ def train_HIPNN_model(model_dir,
               .sum(dim=1)                  # sum along atom axis
              )
 
-            min_dists = hippynn.pretraining.calculate_min_dists(database.arr_dict,species_name=species_key,positions_name=coordinates_key,dist_hard_max=6.0,cell_name=cell_key,batch_size=256)
+            min_dists = hippynn.pretraining.calculate_min_dists(database.arr_dict,species_name=species_key,positions_name=coordinates_key,dist_hard_max=6.0,cell_name=cell_key,batch_size=256,device=device_cuda)
             min_dists_flag = min_dists>network_params["dist_soft_min"]
             for curK in database.arr_dict.keys():
                 database.arr_dict[curK] = database.arr_dict[curK][min_dists_flag]