fix

Author: OutisLi

deepmd/pt/utils/utils.py

@@ -227,10 +227,6 @@ def to_numpy_array(xx: torch.Tensor) -> np.ndarray: ...
 def to_numpy_array(xx: None) -> None: ...
 
 
-@overload
-def to_numpy_array(xx: float) -> np.ndarray: ...
-
-
 def to_numpy_array(
     xx: torch.Tensor | np.ndarray | float | None,
 ) -> np.ndarray | None:

deepmd/tf/utils/learning_rate.py

@@ -13,6 +13,7 @@
     BaseLR,
 )
 from deepmd.tf.env import (
+    GLOBAL_TF_FLOAT_PRECISION,
     tf,
 )
 
@@ -90,13 +91,11 @@ def build(self, global_step: tf.Tensor, num_steps: int) -> tf.Tensor:
         self._base_lr = BaseLR(**params)
 
         # === Step 2. Bind a numpy_function for runtime evaluation ===
-        from deepmd.tf.env import (
-            GLOBAL_TF_FLOAT_PRECISION,
-        )
+        base_lr = self._base_lr
 
         def _lr_value(step: np.ndarray) -> np.ndarray:
             return np.asarray(
-                self._base_lr.value(step),
+                base_lr.value(step),
                 dtype=GLOBAL_TF_FLOAT_PRECISION.as_numpy_dtype,
             )