From 63316d68d5b226ee029c9f1e388e133126996b8c Mon Sep 17 00:00:00 2001
From: Grain Team <grain-dev@google.com>
Date: Fri, 7 Aug 2026 11:26:02 -0700
Subject: [PATCH] internal

PiperOrigin-RevId: 961031826
---
 grain/_src/python/checkpoint/handler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/grain/_src/python/checkpoint/handler.py b/grain/_src/python/checkpoint/handler.py
index 864277e5d..9dd3a7976 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.restore_state(state)  # pyrefly: ignore[bad-argument-type]
     item.start_prefetch()
     return item