diff --git a/sc2/bot_ai.py b/sc2/bot_ai.py
index 19045be78..c85f41356 100644
--- a/sc2/bot_ai.py
+++ b/sc2/bot_ai.py
@@ -148,7 +148,7 @@ async def distribute_workers(self):
             for x in range(0, deficit):
                 if worker_pool:
                     w = worker_pool.pop()
-                    if len(w.orders) == 1 and w.orders[0].ability.id in [AbilityId.HARVESTRETURN]:
+                    if len(w.orders) == 1 and w.orders[0].ability.id in [AbilityId.HARVEST_RETURN]:
                         await self.do(w.move(g))
                         await self.do(w.return_resource(queue=True))
                     else:
@@ -163,7 +163,7 @@ async def distribute_workers(self):
                 if worker_pool:
                     w = worker_pool.pop()
                     mf = self.state.mineral_field.closest_to(townhall)
-                    if len(w.orders) == 1 and w.orders[0].ability.id in [AbilityId.HARVESTRETURN]:
+                    if len(w.orders) == 1 and w.orders[0].ability.id in [AbilityId.HARVEST_RETURN]:
                         await self.do(w.move(townhall))
                         await self.do(w.return_resource(queue=True))
                         await self.do(w.gather(mf, queue=True))