diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py
index 4e5cf13a0..2d50e4496 100644
--- a/langgraph/pregel/__init__.py
+++ b/langgraph/pregel/__init__.py
@@ -482,7 +482,7 @@ class Pregel(
             # apply to checkpoint and save
             _apply_writes(checkpoint, channels, task.writes)
             return self.checkpointer.put(
-                config,
+                saved.config if saved else config,
                 create_checkpoint(checkpoint, channels),
                 {
                     "source": "update",
@@ -551,7 +551,7 @@ class Pregel(
             # apply to checkpoint and save
             _apply_writes(checkpoint, channels, task.writes)
             return await self.checkpointer.aput(
-                config,
+                saved.config if saved else config,
                 create_checkpoint(checkpoint, channels),
                 {
                     "source": "update",