diff --git a/src/qrisp/qaoa/problems/QUBO.py b/src/qrisp/qaoa/problems/QUBO.py
index 5a1d5651..27ac1584 100644
--- a/src/qrisp/qaoa/problems/QUBO.py
+++ b/src/qrisp/qaoa/problems/QUBO.py
@@ -77,15 +77,7 @@ def create_QUBO_cost_operator(Q):
     cost_operator function.
 
     """
-    #def QUBO_cost_operator(qv, gamma):
-        
-    #    for i in range(len(Q)):
-    #        rz(-0.5*2*gamma*(0.5*Q[i][i]+0.5*sum(Q[i])), qv[i])
-    #        for j in range(i+1, len(qv)):
-    #            if Q[i][j] !=0:
-    #                rzz(0.25*2*gamma*Q[i][j], qv[i], qv[j])
-    #return QUBO_cost_operator
-    #new try
+
     def QUBO_cost_operator(qv, gamma):
 
         gphase(-gamma/4*(np.sum(Q)+np.trace(Q)),qv[0])
@@ -119,7 +111,7 @@ def QUBO_problem(Q):
     return QAOAProblem(create_QUBO_cost_operator(Q), RX_mixer, create_QUBO_cl_cost_function(Q))
 
 
-def solve_QUBO(Q, depth, backend = None, n_solutions = 1, print_res = True):
+def solve_QUBO(Q, depth, max_iter = 50, backend = None, n_solutions = 1, print_res = False):
     """
     Solves a Quadratic Unconstrained Binary Optimization (QUBO) problem using the Quantum Approximate Optimization Algorithm (QAOA). 
     The function imports the default backend from the 'qrisp.default_backend' module. 
@@ -134,6 +126,8 @@ def solve_QUBO(Q, depth, backend = None, n_solutions = 1, print_res = True):
         QUBO matrix to solve.
     depth : int
         The depth (amount of layers) of the QAOA circuit.
+    max_iter : int
+        The maximal amount of iterations of the COBYLA optimizer in the QAOA algorithm.
     backend : str
         The backend to be used for the quantum/annealing simulation.
     n_solutions : int
@@ -141,8 +135,8 @@ def solve_QUBO(Q, depth, backend = None, n_solutions = 1, print_res = True):
 
     Returns
     -------
-    None
-        The function prints the runtime of the QAOA algorithm and the ``n_solutions`` best solutions with their respective costs.
+    optimal_solution: tuple
+        The function returns the optimal solution as a tuple where the first element is the cost and the second element is the optimal bitstring. If print_res is set to True, the function prints the runtime of the QAOA algorithm and the ``n_solutions`` best solutions with their respective costs.
 
     """
 
@@ -159,7 +153,7 @@ def solve_QUBO(Q, depth, backend = None, n_solutions = 1, print_res = True):
         backend = backend
 
     # Run QAOA with given quantum arguments, depth, measurement keyword arguments and maximum iterations for optimization
-    res = QUBO_instance.run(qarg, depth, mes_kwargs={"backend" : backend}, max_iter = 50) # runs the simulation
+    res = QUBO_instance.run(qarg, depth, mes_kwargs={"backend" : backend}, max_iter = max_iter) # runs the simulation
     res = dict(list(res.items())[:n_solutions])
 
     # Calculate the cost for each solution
@@ -168,7 +162,12 @@ def solve_QUBO(Q, depth, backend = None, n_solutions = 1, print_res = True):
     # Sort the solutions by their cost in ascending order
     sorted_costs_and_solutions = sorted(costs_and_solutions, key=itemgetter(0))
 
+    optimal_solution = sorted_costs_and_solutions[0]
+
     if print_res is True:
         # Get the top solutions and print them
         for i in range(n_solutions):
-            print(f"Solution {i+1}: {sorted_costs_and_solutions[i][1]} with cost: {sorted_costs_and_solutions[i][0]}")
\ No newline at end of file
+            print(f"Solution {i+1}: {sorted_costs_and_solutions[i][1]} with cost: {sorted_costs_and_solutions[i][0]}")
+    
+    return optimal_solution
+