diff --git a/kadmos/graph/graph_data.py b/kadmos/graph/graph_data.py
index 5f9db82d1c3b5015849135001a3df8c1c9a9bffc..8ba8deb0c747a97283f6d45f7fd48acc885332ab 100644
--- a/kadmos/graph/graph_data.py
+++ b/kadmos/graph/graph_data.py
@@ -3286,7 +3286,7 @@ class FundamentalProblemGraph(DataGraph, KeChainMixin):
         # Add partition id to the nodes
         for idx, partition in enumerate(best_partitions):
             for node in partition:
-                self.nodes[node]['partition_id'] = id
+                self.nodes[node]['partition_id'] = idx
 
         # Add partition to the input graph
         if 'problem_formulation' not in self.graph: