diff --git a/qiskit_optimization/algorithms/grover_optimizer.py b/qiskit_optimization/algorithms/grover_optimizer.py
index 7f73a1bd6..ca14d64dd 100644
--- a/qiskit_optimization/algorithms/grover_optimizer.py
+++ b/qiskit_optimization/algorithms/grover_optimizer.py
@@ -161,8 +161,7 @@ def solve(self, problem: QuadraticProgram) -> OptimizationResult:
         problem_init = deepcopy(problem_)
 
         # convert to minimization problem
-        sense = problem_.objective.sense
-        if sense == problem_.objective.Sense.MAXIMIZE:
+        if problem_.objective.sense == problem_.objective.Sense.MAXIMIZE:
             problem_.objective.sense = problem_.objective.Sense.MINIMIZE
             problem_.objective.constant = -problem_.objective.constant
             for i, val in problem_.objective.linear.to_dict().items():
@@ -176,7 +175,7 @@ def solve(self, problem: QuadraticProgram) -> OptimizationResult:
         optimum_key = math.inf
         optimum_value = math.inf
         threshold = 0
-        n_key = len(problem_.variables)
+        n_key = self._num_key_qubits
         n_value = self._num_value_qubits
 
         # Variables for tracking the solutions encountered.
@@ -236,6 +235,20 @@ def solve(self, problem: QuadraticProgram) -> OptimizationResult:
                     logger.info('Current Optimum Value: %s', optimum_value)
                     improvement_found = True
                     threshold = optimum_value
+
+                    # trace out work qubits and store samples
+                    if self._quantum_instance.is_statevector:  # type: ignore
+                        indices = list(range(n_key, len(outcome)))
+                        rho = partial_trace(self._circuit_results, indices)
+                        self._circuit_results = np.diag(rho.data) ** 0.5
+                    else:
+                        self._circuit_results = {i[0:n_key]: v for i,
+                                                 v in self._circuit_results.items()}
+
+                    raw_samples = self._eigenvector_to_solutions(self._circuit_results,
+                                                                 problem_init)
+                    raw_samples.sort(key=lambda x: problem_.objective.sense.value * x.fval)
+                    samples = self._interpret_samples(problem, raw_samples, self._converters)
                 else:
                     # Using Durr and Hoyer method, increase m.
                     m = int(np.ceil(min(m * 8 / 7, 2 ** (n_key / 2))))
@@ -251,19 +264,6 @@ def solve(self, problem: QuadraticProgram) -> OptimizationResult:
                         improvement_found = True
                         optimum_found = True
 
-                # trace out work qubits
-                if self._quantum_instance.is_statevector:  # type: ignore
-                    indices = list(range(n_key, len(outcome)))
-                    rho = partial_trace(self._circuit_results, indices)
-                    self._circuit_results = np.diag(rho.data) ** 0.5
-                else:
-                    self._circuit_results = {i[0:n_key]: v for i,
-                                             v in self._circuit_results.items()}
-
-                raw_samples = self._eigenvector_to_solutions(self._circuit_results, problem_init)
-                raw_samples.sort(key=lambda x: problem_.objective.sense.value * x.fval)
-                samples = self._interpret_samples(problem, raw_samples, self._converters)
-
                 # Track the operation count.
                 operations = circuit.count_ops()
                 operation_count[iteration] = operations
@@ -313,7 +313,7 @@ def _get_probs(self, qc: QuantumCircuit) -> Dict[str, float]:
             state = result.get_counts(qc)
             shots = self.quantum_instance.run_config.shots
             hist = {key[::-1]: val / shots for key, val in state.items() if val > 0}
-            self._circuit_results = {b[::-1]: (v / shots) ** 0.5 for (b, v) in state.items()}
+            self._circuit_results = {b[::-1]: np.sqrt(v / shots) for (b, v) in state.items()}
         return hist
 
     @staticmethod
diff --git a/test/algorithms/test_grover_optimizer.py b/test/algorithms/test_grover_optimizer.py
index 718f4d0b7..0047a9628 100644
--- a/test/algorithms/test_grover_optimizer.py
+++ b/test/algorithms/test_grover_optimizer.py
@@ -179,7 +179,7 @@ def test_samples_and_raw_samples(self):
         success = OptimizationResultStatus.SUCCESS
         algorithm_globals.random_seed = 1
         grover_optimizer = GroverOptimizer(
-            5, num_iterations=2, quantum_instance=self.qasm_simulator)
+            8, num_iterations=5, quantum_instance=self.qasm_simulator)
         result = grover_optimizer.solve(op)
         self.assertEqual(len(result.samples), 8)
         self.assertEqual(len(result.raw_samples), 32)
@@ -190,9 +190,9 @@ def test_samples_and_raw_samples(self):
         self.assertAlmostEqual(min(s.fval for s in result.raw_samples), opt_sol)
         for sample in result.raw_samples:
             self.assertEqual(sample.status, success)
-        np.testing.assert_array_almost_equal(result.x, result.samples[0].x)
-        self.assertAlmostEqual(result.fval, result.samples[0].fval)
-        self.assertEqual(result.status, result.samples[0].status)
+        np.testing.assert_array_almost_equal(result.x, result.raw_samples[0].x[0:2])
+        self.assertAlmostEqual(result.fval, result.raw_samples[0].fval)
+        self.assertEqual(result.status, result.raw_samples[0].status)
 
 
 if __name__ == '__main__':