diff --git a/flatland/envs/rail_generators.py b/flatland/envs/rail_generators.py
index de7804024dba7540786c90f9537f9babcede206b..147b1bd14a5943f87f4d78f37e9ea7e1722c9087 100644
--- a/flatland/envs/rail_generators.py
+++ b/flatland/envs/rail_generators.py
@@ -781,7 +781,7 @@ def sparse_rail_generator(num_cities=5, num_intersections=4, num_trainstations=2
             node_available_target.append(len(train_stations[node_idx]))
 
         # Assign agents to slots
-        for agent_idx in range(nb_nodes):
+        for agent_idx in range(num_agents):
             avail_start_nodes = [idx for idx, val in enumerate(node_available_start) if val > 0]
             avail_target_nodes = [idx for idx, val in enumerate(node_available_target) if val > 0]
             start_node = np.random.choice(avail_start_nodes)
diff --git a/flatland/envs/schedule_generators.py b/flatland/envs/schedule_generators.py
index b3576a2bec77f75afc9331cc6c190649590a990c..8c1834901a1e132d6adfa54cdddb286aafbef2a7 100644
--- a/flatland/envs/schedule_generators.py
+++ b/flatland/envs/schedule_generators.py
@@ -68,6 +68,7 @@ def sparse_schedule_generator(speed_ratio_map: Mapping[float, float] = None) ->
         agents_position = []
         agents_target = []
         agents_direction = []
+
         for agent_idx in range(num_agents):
             # Set target for agent
             current_target_node = agent_start_targets_nodes[agent_idx][1]