From 238c940131aaf3e1b66e69aac78a81cc8d468bd2 Mon Sep 17 00:00:00 2001
From: shanyu <santohyunolfsssigneb@gmail.com>
Date: Sat, 26 Sep 2026 02:16:57 +0800
Subject: [PATCH] Remember the first PID gradient for the next derivative

The first step assigned the local name g_buf to the current gradient
after storing a zero buffer. The stored gradient stayed 0, and the
next step took its derivative against that zero instead of the
gradient from the previous step.

Copy the gradient into the buffer. The first derivative is still 0,
because there is no earlier gradient. The second step then uses
(1 - momentum) times the change in gradient.
---
 tests/test_pid_grad_buffer.py | 28 ++++++++++++++++++++++++++++
 torch_optimizer/pid.py        |  5 ++++-
 2 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_pid_grad_buffer.py

diff --git a/tests/test_pid_grad_buffer.py b/tests/test_pid_grad_buffer.py
new file mode 100644
index 00000000..ecdbdd2f
--- /dev/null
+++ b/tests/test_pid_grad_buffer.py
@@ -0,0 +1,28 @@
+"""The first PID step must remember its gradient."""
+import torch
+import torch_optimizer as optim
+
+
+def test_second_step_uses_the_previous_gradient():
+    param = torch.nn.Parameter(torch.tensor([0.0]))
+    opt = optim.PID(
+        [param],
+        lr=1.0,
+        momentum=0.5,
+        dampening=0.0,
+        weight_decay=0.0,
+        integral=0.0,
+        derivative=10.0,
+    )
+    param.grad = torch.tensor([1.0])
+    opt.step()
+    assert torch.equal(
+        opt.state[param]["grad_buffer"], torch.tensor([1.0])
+    )
+    # first derivative is 0: there is no earlier gradient
+    assert torch.equal(param.detach(), torch.tensor([-1.0]))
+
+    param.grad = torch.tensor([3.0])
+    opt.step()
+    # D = (1 - momentum) * (3 - 1) = 1, so the step is 3 + 10 * 1
+    assert torch.allclose(param.detach(), torch.tensor([-14.0]))
diff --git a/torch_optimizer/pid.py b/torch_optimizer/pid.py
index 44dad253..ca6bf1c4 100644
--- a/torch_optimizer/pid.py
+++ b/torch_optimizer/pid.py
@@ -104,7 +104,10 @@ def step(self, closure: OptLossClosure = None) -> OptFloat:
                         g_buf = param_state["grad_buffer"] = torch.zeros_like(
                             p, memory_format=torch.preserve_format
                         )
-                        g_buf = d_p
+                        # copy into the buffer. Rebinding g_buf = d_p left
+                        # the stored gradient at 0, so the next step's
+                        # derivative was taken against a zero previous grad.
+                        g_buf.copy_(d_p)
 
                         d_buf = param_state["d_buffer"] = torch.zeros_like(
                             p, memory_format=torch.preserve_format