How to efficiently create a large multilayer graph from ROOT file for machine learning uses in StellarGraph

Question

I want to create a multilayered graph following this kind of structure (but with every node also connecting to each node from next layer):

What I want to do with this graph is to apply some machine learning algorithm after loading it to StellarGraph, so I'm using code from NetworkNx to try to create the graph since StellarGraph can load graphs directly from NetworkNx.

Since creating the whole graph at once would be too hard, I tried creating the graph by batches, heres some code for creating a graph for the layer 1:

file = uproot.open("file.root")
tree = file["data"]


# Get the Event branches
hits = tree["Nhit"].array()
i_pos_arrays = tree["I"].array()
j_pos_arrays = tree["J"].array()
layers = tree["K"].array()

node_index = 1
graph = nx.Graph()
previous = []
layer = 1
start_time = time.time()

for hits, x, y, z in zip(hits, i_pos_arrays,j_pos_arrays,layers):
   sorted_indices = np.argsort(z)
   x = x[sorted_indices]
   y = y[sorted_indices]
   z = z[sorted_indices]

   for i in range(hits):
        if(z[i] == layer):
            #node_attributes = {'x': x[i], 'y': y[i] ,'z':z[i]}
            graph.add_node(node_index)
           # graph.nodes[node_index].update(node_attributes)  # Luego actualizar atributos

            if previous:
                for pre_graphs in previous:
                    graph.add_edge(pre_graphs, node_index)
                previous.append(node_index)
            else:
                previous.append(node_index)
            node_index = node_index + 1

Info: every element in i_pos_arrays, j_pos_arrays, layers are arrays, hits is a single array

The problem: I simply get out of memory while creating the graph, (16GB RAM), so I want to know what other tools or methods I can use to reach my goal.