fix

Author: OutisLi

deepmd/tf/utils/learning_rate.py

@@ -13,6 +13,7 @@
     BaseLR,
 )
 from deepmd.tf.env import (
+    GLOBAL_TF_FLOAT_PRECISION,
     tf,
 )
 
@@ -93,14 +94,15 @@ def build(self, global_step: tf.Tensor, num_steps: int) -> tf.Tensor:
         base_lr = self._base_lr
 
         def _lr_value(step: np.ndarray) -> np.ndarray:
-            # Use float32 for learning rate, consistent with PyTorch/Paddle backends
+            # Use GLOBAL_TF_FLOAT_PRECISION (float64) for learning rate,
+            # consistent with energy precision in TF backend
             return np.asarray(
                 base_lr.value(step),
-                dtype=np.float32,
+                dtype=GLOBAL_TF_FLOAT_PRECISION.as_numpy_dtype,
             )
 
         lr = tf.numpy_function(
-            _lr_value, [global_step], Tout=tf.float32, name="lr_schedule"
+            _lr_value, [global_step], Tout=GLOBAL_TF_FLOAT_PRECISION, name="lr_schedule"
         )
         lr.set_shape(global_step.get_shape())
         return lr

source/tests/tf/test_lr.py

@@ -42,13 +42,13 @@ class TestLearningRateScheduleBuild(unittest.TestCase):
     """Test TF tensor building and integration."""
 
     def test_build_returns_tensor(self) -> None:
-        """Test that build() returns a float32 TF tensor (consistent with PT/PD backends)."""
+        """Test that build() returns a float64 TF tensor (consistent with GLOBAL_TF_FLOAT_PRECISION)."""
         lr_schedule = LearningRateSchedule({"start_lr": 1e-3, "stop_lr": 1e-5})
         global_step = tf.constant(0, dtype=tf.int64)
         lr_tensor = lr_schedule.build(global_step, num_steps=10000)
 
         self.assertIsInstance(lr_tensor, tf.Tensor)
-        self.assertEqual(lr_tensor.dtype, tf.float32)
+        self.assertEqual(lr_tensor.dtype, tf.float64)
 
     def test_default_type_exp(self) -> None:
         """Test that default type is 'exp' when not specified."""