diff --git a/benchmark.py b/benchmark.py
index 74da33b..f5c6a89 100644
--- a/benchmark.py
+++ b/benchmark.py
@@ -47,7 +47,7 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..',
 from hdh import HDH, plot_hdh
 from hdh.converters.qiskit import from_qiskit
 from hdh.converters.qasm import from_qasm
-from hdh.passes.cut import compute_cut, cost, partition_sizes, compute_parallelism_by_time
+from hdh.passes.cut import compute_cut, cost, partition_size, compute_parallelism_by_time
 
 # Circuit examples
 from circuit_examples import HDHCircuitLibrary
@@ -167,7 +167,8 @@ class HDHBenchmarkSuite:
                     start_partition = time.perf_counter()
                     partitions, _ = compute_cut(hdh, num_parts, capacity)
                     partitioning_time = time.perf_counter() - start_partition
-                    partition_cost = cost(hdh, partitions)
+                    cost_q, cost_c = cost(hdh, partitions)
+                    partition_cost = cost_q + cost_c  # Total cost
                 except Exception as e:
                     self.logger.warning(f"Partitioning failed for {circuit_name}: {str(e)}")
             
diff --git a/main.py b/main.py
index 4349be6..74872c8 100644
--- a/main.py
+++ b/main.py
@@ -36,7 +36,7 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..',
 from hdh import HDH, plot_hdh
 from hdh.converters.qiskit import from_qiskit
 from hdh.converters.qasm import from_qasm
-from hdh.passes.cut import compute_cut, cost, partition_sizes, compute_parallelism_by_time
+from hdh.passes.cut import compute_cut, cost, partition_size, compute_parallelism_by_time
 
 # Quantum computing imports
 from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister