Hack AC to not clear recomputed activations

soulitzer · soulitzer · commit 1a1057279e77 · 2025-02-03T12:29:37.000-05:00
ghstack-source-id: f9c5294 Pull Request resolved: #146318
diff --git a/torch/utils/checkpoint.py b/torch/utils/checkpoint.py
@@ -1108,7 +1108,7 @@ def pack_hook(x):
             return holder
 
         def unpack_hook(holder):
-            gid = torch._C._current_graph_task_id()
+            gid = 0 # torch._C._current_graph_task_id()
             if gid == -1:
                 # generate a temporary id if we trigger unpack outside of a backward call
                 gid = int(uuid.uuid4())
@@ -1137,7 +1137,7 @@ def unpack_hook(holder):
                 )
             _internal_assert(holder.handles[gid] in frame.recomputed[gid])
             ret = frame.recomputed[gid][holder.handles[gid]]
-            holder.handles[gid] = None
+            # holder.handles[gid] = None
             return ret
 
         if frame.unpack_error_cb is not None:
