diff --git a/tests/test_flatland_envs_predictions.py b/tests/test_flatland_envs_predictions.py
index 45e0bdda32fdbf4135c9400d655e49f146bce08c..cfb3e9826d3d2d262a3e28a637f5012eabeed2b8 100644
--- a/tests/test_flatland_envs_predictions.py
+++ b/tests/test_flatland_envs_predictions.py
@@ -34,7 +34,7 @@ def test_dummy_predictor(rendering=False):
 
     # reset to set agents from agents_static
     env.reset(False, False)
-    env.set_agent_active(0)
+    env.set_agent_active(env.agents[0])
 
     if rendering:
         renderer = RenderTool(env, gl="PILSVG")