diff --git a/simdkalman/kalmanfilter.py b/simdkalman/kalmanfilter.py
index 2c723de..978a0af 100644
--- a/simdkalman/kalmanfilter.py
+++ b/simdkalman/kalmanfilter.py
@@ -519,7 +519,7 @@ def auto_flat_states(obs_gaussian):
             filtered_states = empty_gaussian(cov=True)
 
         if gains:
-            result.filtered.gains = np.empty((n_vars, n_measurements, n_states, n_states))
+            result.filtered.gains = np.empty((n_vars, n_measurements, n_states, n_obs))
 
         for j in range(n_measurements):
             if verbose: