diff --git a/cognee/modules/graph/utils/expand_with_nodes_and_edges.py b/cognee/modules/graph/utils/expand_with_nodes_and_edges.py
index 775930638..e1f61d115 100644
--- a/cognee/modules/graph/utils/expand_with_nodes_and_edges.py
+++ b/cognee/modules/graph/utils/expand_with_nodes_and_edges.py
@@ -144,6 +144,7 @@ def expand_with_nodes_and_edges(
                     is_a=type_node,
                     description=node.description,
                     ontology_valid=ontology_validated_source_ent,
+                    belongs_to_set=data_chunk.belongs_to_set,
                 )
 
                 added_nodes_map[entity_node_key] = entity_node