diff --git a/python/quadruped_reactive_walking/WB_MPC_Wrapper.py b/python/quadruped_reactive_walking/WB_MPC_Wrapper.py
index 9de9dff2d4dfa9ab5c210303d7fcdfb371df3be8..72008a4b0ce51be400337853f3095e49f0c7ffec 100644
--- a/python/quadruped_reactive_walking/WB_MPC_Wrapper.py
+++ b/python/quadruped_reactive_walking/WB_MPC_Wrapper.py
@@ -15,7 +15,7 @@ class Result:
     def __init__(self, pd):
         self.xs = list(np.zeros((pd.T + 1, pd.nx)))
         self.us = list(np.zeros((pd.T, pd.nu)))
-        self.K = list(np.zeros([pd.T, pd.nu, pd.nx]))
+        self.K = list(np.zeros([pd.T, pd.nu, pd.ndx]))
         self.solving_duration = 0.0
         self.new_result = False
 
@@ -47,7 +47,7 @@ class MPC_Wrapper:
             self.in_gait = Array("d", [0] * (pd.T * 4))
             self.out_xs = Array("d", [0] * ((pd.T + 1) * pd.nx))
             self.out_us = Array("d", [0] * (pd.T * pd.nu))
-            self.out_k = Array("d", [0] * (pd.T * pd.nu * pd.nx))
+            self.out_k = Array("d", [0] * (pd.T * pd.nu * pd.ndx))
             self.out_solving_time = Value("d", 0.0)
         else:
             self.ocp = OCP(pd, footsteps, gait)
@@ -215,7 +215,7 @@ class MPC_Wrapper:
             ] = np.array(us)
         with self.out_k.get_lock():
             np.frombuffer(self.out_k.get_obj()).reshape(
-                [self.pd.T, self.pd.nu, self.pd.nx]
+                [self.pd.T, self.pd.nu, self.pd.ndx]
             )[:, :, :] = np.array(K)
         self.out_solving_time.value = solving_time
 
@@ -230,7 +230,7 @@ class MPC_Wrapper:
         us = list(np.frombuffer(self.out_us.get_obj()).reshape((self.pd.T, self.pd.nu)))
         K = list(
             np.frombuffer(self.out_k.get_obj()).reshape(
-                [self.pd.T, self.pd.nu, self.pd.nx]
+                [self.pd.T, self.pd.nu, self.pd.ndx]
             )
         )
         solving_time = self.out_solving_time.value