diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py
index f309a46fa..b63295c92 100644
--- a/mattport/nerf/graph.py
+++ b/mattport/nerf/graph.py
@@ -19,7 +19,7 @@ class Node:
     """Node datastructure for graph composition."""
 
     name: str
-    children: set
+    children: dict
 
     def __hash__(self):
         return hash(self.name)
@@ -44,7 +44,7 @@ class Graph(nn.Module):
         # generate dependency ordering for module calls
         self.module_order = self.get_module_order()
 
-    def construct_graph(self) -> set:
+    def construct_graph(self) -> dict:
         """Constructs a dependency graph given the module configuration
 
         Args:
@@ -57,17 +57,17 @@ class Graph(nn.Module):
         roots = set()
         for module_name, module_dict in self.modules_config.items():
             if not module_name in processed_modules:
-                curr_module = Node(name=module_name, children=set())
+                curr_module = Node(name=module_name, children={})
                 processed_modules[module_name] = curr_module
             else:
                 curr_module = processed_modules[module_name]
             inputs = module_dict["inputs"]
             for input_module in inputs:
                 if not input_module in processed_modules:
-                    parent_module = Node(name=input_module, children=set([curr_module]))
+                    parent_module = Node(name=input_module, children={module_name: curr_module})
                     processed_modules[input_module] = parent_module
                 else:
-                    processed_modules[input_module].children.add(curr_module)
+                    processed_modules[input_module].children[module_name] = curr_module
                 if input_module == "x":
                     roots.add(curr_module)
         return roots
diff --git a/tests/nerf/test_graph.py b/tests/nerf/test_graph.py
index 96a10d1ce..9a0d20260 100644
--- a/tests/nerf/test_graph.py
+++ b/tests/nerf/test_graph.py
@@ -89,12 +89,12 @@ def create_complex_graph():
 def check_consistency(root: "Node", targets: list):
     """rudimentary check on predicted graph versus expected target ordering"""
     curr_pointer = root
-    for target_name, target_children in targets:
+    for i, (target_name, target_children) in enumerate(targets):
         assert curr_pointer.name == target_name
         assert len(curr_pointer.children) == target_children
         if len(curr_pointer.children) > 0:
-            # for now traversing smaller parts of graph.. need stabler checks
-            curr_pointer = list(curr_pointer.children)[0]
+            assert targets[i + 1][0] in curr_pointer.children
+            curr_pointer = curr_pointer.children[targets[i + 1][0]]
 
 
 def test_graph_init():
@@ -146,5 +146,14 @@ def test_construct_graph_complex():
     roots = test_graph.construct_graph()
     assert len(roots) == 1
     roots = list(roots)
-    targets = [("encoder_0", 3), ("mlp_0", 2), ("mlp_3", 0)]
-    check_consistency(roots[0], targets)
+    targets_p0 = [("encoder_0", 3), ("mlp_0", 2), ("mlp_3", 0)]
+    check_consistency(roots[0], targets_p0)
+
+    targets_p1 = [("encoder_0", 3), ("mlp_0", 2), ("mlp_2", 0)]
+    check_consistency(roots[0], targets_p1)
+
+    targets_p2 = [("encoder_0", 3), ("mlp_1", 0)]
+    check_consistency(roots[0], targets_p2)
+
+    targets_p3 = [("encoder_0", 3), ("mlp_2", 0)]
+    check_consistency(roots[0], targets_p3)