diff --git a/deepmd/utils/learning_rate.py b/deepmd/utils/learning_rate.py
index 0f68f2a7df..5997b7d63a 100644
--- a/deepmd/utils/learning_rate.py
+++ b/deepmd/utils/learning_rate.py
@@ -45,7 +45,7 @@ def __init__(
             self.decay_rate = decay_rate
         self.min_lr = stop_lr
 
-    def value(self, step) -> np.ndarray:
+    def value(self, step) -> np.float64:
         """Get the learning rate at the given step."""
         step_lr = self.start_lr * np.power(self.decay_rate, step // self.decay_steps)
         if step_lr < self.min_lr: