diff --git a/README.md b/README.md
index 626d93f1bb..100dcec4c2 100644
--- a/README.md
+++ b/README.md
@@ -19,7 +19,7 @@ For more information, check the [documentation](https://deepmd.readthedocs.io/).
 
 ### Highlighted features
 
-- **interfaced with multiple backends**, including TensorFlow, PyTorch, JAX and Paddle the most popular deep learning frameworks, making the training process highly automatic and efficient.
+- **interfaced with multiple backends**, including TensorFlow, PyTorch, JAX and Paddle, the most popular deep learning frameworks, making the training process highly automatic and efficient.
 - **interfaced with high-performance classical MD and quantum (path-integral) MD packages**, including LAMMPS, i-PI, AMBER, CP2K, GROMACS, OpenMM, and ABUCUS.
 - **implements the Deep Potential series models**, which have been successfully applied to finite and extended systems, including organic molecules, metals, semiconductors, insulators, etc.
 - **implements MPI and GPU supports**, making it highly efficient for high-performance parallel and distributed computing.
diff --git a/deepmd/pd/optimizer/KFWrapper.py b/deepmd/pd/optimizer/KFWrapper.py
index 7fd8506300..32838fc389 100644
--- a/deepmd/pd/optimizer/KFWrapper.py
+++ b/deepmd/pd/optimizer/KFWrapper.py
@@ -58,7 +58,8 @@ def update_energy(
         mask = error < 0
 
         error = error * update_prefactor
-        error[mask] = -1 * error[mask]
+        # error[mask] = -1 * error[mask]
+        error = _mask_update(error, mask, -error[mask])
         error = error.mean()
 
         if self.is_distributed:
@@ -66,7 +67,8 @@ def update_energy(
             error /= dist.get_world_size()
 
         Etot_predict = update_prefactor * Etot_predict
-        Etot_predict[mask] = -Etot_predict[mask]
+        # Etot_predict[mask] = -Etot_predict[mask]
+        Etot_predict = _mask_update(Etot_predict, mask, -Etot_predict[mask])
 
         Etot_predict.sum().backward()
         error = error * math.sqrt(bs)
@@ -91,7 +93,7 @@ def update_force(
             error_tmp = Force_label[:, index[i]] - force_predict[:, index[i]]
             error_tmp = update_prefactor * error_tmp
             mask = error_tmp < 0
-            error_tmp = _mask_update(error_tmp, mask, -1 * error_tmp[mask])
+            error_tmp = _mask_update(error_tmp, mask, -error_tmp[mask])
             # error_tmp[mask] = -1 * error_tmp[mask]
             error = error_tmp.mean() / natoms_sum
 
diff --git a/deepmd/pd/optimizer/LKF.py b/deepmd/pd/optimizer/LKF.py
index 06e4e2e156..d77b4a9232 100644
--- a/deepmd/pd/optimizer/LKF.py
+++ b/deepmd/pd/optimizer/LKF.py
@@ -265,6 +265,7 @@ def __update(self, H, error, weights):
     def set_grad_prefactor(self, grad_prefactor):
         self.grad_prefactor = grad_prefactor
 
+    @paddle.no_grad()
     def step(self, error):
         params_packed_index = self._state.get("params_packed_index")
 
diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py
index 17b9319f5b..f0d11a4a81 100644
--- a/deepmd/pd/train/training.py
+++ b/deepmd/pd/train/training.py
@@ -42,8 +42,9 @@
     get_model,
     get_zbl_model,
 )
-from deepmd.pd.optimizer import (  # LKFOptimizer,
+from deepmd.pd.optimizer import (
     KFOptimizerWrapper,
+    LKFOptimizer,
 )
 from deepmd.pd.train.wrapper import (
     ModelWrapper,
@@ -601,10 +602,12 @@ def warm_up_linear(step, warmup_steps):
             if optimizer_state_dict is not None and self.restart_training:
                 self.optimizer.set_state_dict(optimizer_state_dict)
         elif self.opt_type == "LKF":
-            raise NotImplementedError("LKF is not supported yet in Paddle backend.")
-            # self.optimizer = LKFOptimizer(
-            #     [{'params': self.wrapper.parameters()}], 0.98, 0.99870, self.opt_param["kf_blocksize"]
-            # )
+            self.optimizer = LKFOptimizer(
+                [{"params": self.wrapper.parameters()}],
+                0.98,
+                0.99870,
+                self.opt_param["kf_blocksize"],
+            )
         else:
             raise ValueError(f"Not supported optimizer type '{self.opt_type}'")
 
diff --git a/source/tests/pd/test_LKF.py b/source/tests/pd/test_LKF.py
index ae9508c149..81f69041da 100644
--- a/source/tests/pd/test_LKF.py
+++ b/source/tests/pd/test_LKF.py
@@ -11,7 +11,6 @@
 )
 
 
-@unittest.skip("Paddle do not support LKF now")
 class TestLKF(unittest.TestCase):
     def test_lkf(self):
         with open(str(Path(__file__).parent / "water/lkf.json")) as fin: