diff --git a/src/deep_neurographs/machine_learning/feature_generation.py b/src/deep_neurographs/machine_learning/feature_generation.py
index 5c45eea..0e17315 100644
--- a/src/deep_neurographs/machine_learning/feature_generation.py
+++ b/src/deep_neurographs/machine_learning/feature_generation.py
@@ -29,10 +29,8 @@
 
 CHUNK_SIZE = [64, 64, 64]
 N_BRANCH_PTS = 50
-N_PROFILE_PTS = 10
+N_PROFILE_PTS = 16
 N_SKEL_FEATURES = 22
-NODE_PROFILE_DEPTH = 15
-WINDOW = [5, 5, 5]
 
 
 def run(
@@ -41,7 +39,7 @@ def run(
     model_type,
     proposals,
     radius,
-    downsample_factor=0,
+    downsample_factor=1,
     labels=None,
 ):
     """
@@ -602,7 +600,7 @@ def get_chunk(img, labels, voxel_1, voxel_2, thread_id=None):
 
     # Generate features
     path = geometry.make_line(patch_voxel_1, patch_voxel_2, N_PROFILE_PTS)
-    profile = geometry.get_profile(chunk, path, window=WINDOW)
+    profile = geometry.get_profile(chunk, path)
     labels_chunk[labels_chunk > 0] = 1
     labels_chunk = geometry.fill_path(labels_chunk, path, val=2)
     chunk = np.stack([chunk, labels_chunk], axis=0)
diff --git a/src/deep_neurographs/machine_learning/heterograph_feature_generation.py b/src/deep_neurographs/machine_learning/heterograph_feature_generation.py
index 0da3a63..1ab1f57 100644
--- a/src/deep_neurographs/machine_learning/heterograph_feature_generation.py
+++ b/src/deep_neurographs/machine_learning/heterograph_feature_generation.py
@@ -17,7 +17,7 @@
 from deep_neurographs.utils import img_util
 
 WINDOW = [5, 5, 5]
-N_PROFILE_PTS = 10
+N_PROFILE_PTS = 16
 NODE_PROFILE_DEPTH = 16
 
 
@@ -82,18 +82,11 @@ def run_on_nodes(neurograph, computation_graph, img, downsample_factor):
     Returns
     -------
     dict
-        Dictionary whose keys are feature types (i.e. skeletal and profiles)
-        and values are a dictionary that maps a node id to the corresponding
-        feature vector.
+        Dictionary whose keys are feature types (i.e. skeletal) and values are
+        a dictionary that maps a node id to the corresponding feature vector.
 
     """
-    node_features = {
-        "skel": node_skeletal(neurograph, computation_graph),
-        "profiles": node_profiles(
-            neurograph, computation_graph, img, downsample_factor
-        ),
-    }
-    return node_features
+    return {"skel": node_skeletal(neurograph, computation_graph)}
 
 
 def run_on_edges(neurograph, computation_graph):
diff --git a/src/deep_neurographs/machine_learning/hetero_graph_models.py b/src/deep_neurographs/machine_learning/heterograph_models.py
similarity index 100%
rename from src/deep_neurographs/machine_learning/hetero_graph_models.py
rename to src/deep_neurographs/machine_learning/heterograph_models.py
diff --git a/src/deep_neurographs/machine_learning/heterograph_trainer.py b/src/deep_neurographs/machine_learning/heterograph_trainer.py
index f0a2230..3554932 100644
--- a/src/deep_neurographs/machine_learning/heterograph_trainer.py
+++ b/src/deep_neurographs/machine_learning/heterograph_trainer.py
@@ -129,7 +129,6 @@ def run_on_graphs(self, datasets, augment=False):
             y, hat_y = [], []
             self.model.train()
             for graph_id in train_ids:
-                print(graph_id)
                 y_i, hat_y_i = self.train(
                     datasets[graph_id], epoch, augment=augment
                 )
diff --git a/src/deep_neurographs/neurograph.py b/src/deep_neurographs/neurograph.py
index 43aa54a..fe25b9b 100644
--- a/src/deep_neurographs/neurograph.py
+++ b/src/deep_neurographs/neurograph.py
@@ -279,10 +279,11 @@ def generate_proposals(
         self,
         search_radius,
         complex_bool=False,
+        groundtruth_graph=None,
         long_range_bool=False,
         proposals_per_leaf=3,
         return_trimmed_proposals=False,
-        trim_endpoints_bool=False,
+        trim_endpoints_bool=False,        
     ):
         """
         Generates proposals from leaf nodes in "self".
@@ -294,6 +295,8 @@ def generate_proposals(
         complex_bool : bool, optional
             Indication of whether to generate complex proposals. The default
             is False.
+        groundtruth_graph : networkx.Graph, optional
+            Ground truth graph. The default is None.
         long_range_bool : bool, optional
             Indication of whether to generate long range proposals. The
             default is False.
@@ -321,6 +324,7 @@ def generate_proposals(
             long_range_bool=long_range_bool,
             trim_endpoints_bool=trim_endpoints_bool,
         )
+        self.target_edges = init_targets(self, groundtruth_graph)
 
     def reset_proposals(self):
         """
@@ -465,9 +469,6 @@ def list_proposals(self):
         """
         return list(self.proposals)
 
-    def init_targets(self, target_neurograph):
-        self.target_edges = init_targets(self, target_neurograph)
-
     # -- KDTree --
     def init_kdtree(self, node_type):
         """