diff --git a/graph_net/subgraph_range_util.py b/graph_net/subgraph_range_util.py
index 6f8ad2be1..97f36a21c 100644
--- a/graph_net/subgraph_range_util.py
+++ b/graph_net/subgraph_range_util.py
@@ -14,7 +14,7 @@ def get_belonging_ranges(rel_model_path, ranges):
             return
         for parent_rel_model_path, parent_ranges in parent_subgraph_ranges[
             rel_model_path
-        ]:
+        ].items():
             yield parent_rel_model_path, flatmap_subgraph_ranges(ranges, parent_ranges)
 
     ret_rel_model_path_and_ranges = [