From ccdfa2fead295d5af7958d7f570b91c5de71cb6b Mon Sep 17 00:00:00 2001
From: Jun Tian <find_my_way@foxmail.com>
Date: Sun, 1 Feb 2026 12:43:21 +0800
Subject: [PATCH] Call `start_prefetch` immediately at the end of restoring

---
 grain/_src/python/checkpoint/handler.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/grain/_src/python/checkpoint/handler.py b/grain/_src/python/checkpoint/handler.py
index d6bebb16b..8eeb2daed 100644
--- a/grain/_src/python/checkpoint/handler.py
+++ b/grain/_src/python/checkpoint/handler.py
@@ -66,6 +66,7 @@ def restore(
     else:
       state = state.encode()
     item.set_state(state)
+    item.start_prefetch()
     return item
 
   # Required by interface but not supported by PyGrain checkpoints.