From 617ecc1a7b4fc4a2682cf5c12dc17631450a85b9 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Tue, 27 May 2025 20:46:55 +0300
Subject: [PATCH] fixup grouper process replay [pr] (#10538)

---
 tinygrad/engine/grouper.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tinygrad/engine/grouper.py b/tinygrad/engine/grouper.py
index 66ce5d2628..4fa9c0a25a 100644
--- a/tinygrad/engine/grouper.py
+++ b/tinygrad/engine/grouper.py
@@ -580,7 +580,7 @@ def get_kernelize_map(big_sink:UOp) -> dict[UOp, UOp]:
     with Context(PICKLE_BUFFERS=0):
       import pickle
       kernel_asts = [u.arg.ast for u in sched_sink.toposort() if u.op is Ops.KERNEL]
-      PROCESS_REPLAY_CAPTURE[id(big_sink)] = pickle.dumps((big_sink, ContextVar._cache, [u.arg.ast for u in kernel_asts if u.op is Ops.KERNEL]))
+      PROCESS_REPLAY_CAPTURE[id(big_sink)] = pickle.dumps((big_sink, ContextVar._cache, kernel_asts))
 
   # map tensors to buffer/assign/const
   # TODO: this is not right, and causes TestDataset.test_dataset_is_realized to fail unless I unprincipledly add Ops.COPY, which breaks others