Update deepmd/utils/learning_rate.py

Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Oct 15, 2024 · 90c6f9a · 90c6f9a
1 parent 367d33e
commit 90c6f9a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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: