diff --git a/tests/test_flatland_malfunction.py b/tests/test_flatland_malfunction.py
index 1727dba83b1b9f1315c2bd09160c4097030a9a9c..b6e082e362eec1185c54bf088c8e54fdef586a56 100644
--- a/tests/test_flatland_malfunction.py
+++ b/tests/test_flatland_malfunction.py
@@ -53,7 +53,7 @@ class SingleAgentNavigationObs(TreeObsForRailEnv):
                     min_distances.append(np.inf)
 
             observation = [0, 0, 0]
-            observation[np.argmin(min_distances)[0]] = 1
+            observation[np.argmin(min_distances)] = 1
 
         return observation