diff --git a/prefect_ray/task_runners.py b/prefect_ray/task_runners.py
index 39a79a1..b73d0fd 100644
--- a/prefect_ray/task_runners.py
+++ b/prefect_ray/task_runners.py
@@ -78,7 +78,7 @@ def count_to(highest_number):
 import anyio
 import ray
 from prefect.futures import PrefectFuture
-from prefect.server.schemas.states import State
+from prefect.client.schemas.objects import State
 from prefect.states import exception_to_crashed_state
 from prefect.task_runners import BaseTaskRunner, R, TaskConcurrencyType
 from prefect.utilities.asyncutils import sync_compatible