diff --git a/grain/_src/python/checkpoint/handler.py b/grain/_src/python/checkpoint/handler.py
index 864277e5d..59522e101 100644
--- a/grain/_src/python/checkpoint/handler.py
+++ b/grain/_src/python/checkpoint/handler.py
@@ -161,7 +161,7 @@ def restore(
       state = json.loads(state)
     else:
       state = state.encode()
-    item.set_state(state)  # pyrefly: ignore[bad-argument-type]
+      item.set_state(state)  # pyrefly: ignore[bad-argument-type]
     item.start_prefetch()
     return item