fix

Author: OutisLi

deepmd/dpmodel/utils/learning_rate.py

@@ -105,12 +105,14 @@ def __init__(
             self.warmup_steps = warmup_steps
 
         # === Step 5. Validate step ranges (runtime check) ===
-        if num_steps <= 0:
-            raise ValueError("num_steps must be positive")
+        if num_steps < 0:
+            raise ValueError("num_steps must be non-negative")
         if self.warmup_steps < 0:
             raise ValueError("warmup_steps must be non-negative")
-        if self.warmup_steps >= num_steps:
+        if num_steps > 0 and self.warmup_steps >= num_steps:
             raise ValueError("warmup_steps must be smaller than num_steps")
+        if num_steps == 0 and self.warmup_steps != 0:
+            raise ValueError("warmup_steps must be 0 when num_steps is 0")
 
         # === Step 6. Compute warmup_start_lr ===
         self.warmup_start_lr = warmup_start_factor * start_lr
@@ -457,6 +459,9 @@ def _decay_value(self, step: int | Array) -> Array:
             step = np.asarray(step)
         xp = array_api_compat.array_namespace(step)
         min_lr = self.start_lr * self.lr_min_factor
+        # Handle decay_num_steps=0 (no training steps) - return start_lr
+        if self.decay_num_steps == 0:
+            return xp.full_like(step, self.start_lr, dtype=xp.float64)
         step_lr = self.start_lr * (
             self.lr_min_factor
             + 0.5

deepmd/tf/loss/dos.py

@@ -65,8 +65,8 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
         dos_hat = label_dict["dos"]
         atom_dos_hat = label_dict["atom_dos"]
 
-        find_dos = label_dict["find_dos"]
-        find_atom_dos = label_dict["find_atom_dos"]
+        find_dos = global_cvt_2_tf_float(label_dict["find_dos"])
+        find_atom_dos = global_cvt_2_tf_float(label_dict["find_atom_dos"])
 
         dos_reshape = tf.reshape(dos, [-1, self.numb_dos])
         dos_hat_reshape = tf.reshape(dos_hat, [-1, self.numb_dos])

deepmd/tf/loss/ener.py

@@ -149,14 +149,14 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
         virial_hat = label_dict["virial"]
         atom_ener_hat = label_dict["atom_ener"]
         atom_pref = label_dict["atom_pref"]
-        find_energy = label_dict["find_energy"]
-        find_force = label_dict["find_force"]
-        find_virial = label_dict["find_virial"]
-        find_atom_ener = label_dict["find_atom_ener"]
-        find_atom_pref = label_dict["find_atom_pref"]
+        find_energy = global_cvt_2_ener_float(label_dict["find_energy"])
+        find_force = global_cvt_2_tf_float(label_dict["find_force"])
+        find_virial = global_cvt_2_tf_float(label_dict["find_virial"])
+        find_atom_ener = global_cvt_2_tf_float(label_dict["find_atom_ener"])
+        find_atom_pref = global_cvt_2_tf_float(label_dict["find_atom_pref"])
         if self.has_gf:
             drdq = label_dict["drdq"]
-            find_drdq = label_dict["find_drdq"]
+            find_drdq = global_cvt_2_tf_float(label_dict["find_drdq"])
         else:
             find_drdq = 0.0
 
@@ -589,10 +589,10 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
         virial_label = label_dict["virial"]
         atom_ener_label = label_dict["atom_ener"]
         atom_pref = label_dict["atom_pref"]
-        find_energy = label_dict["find_energy"]
-        find_force = label_dict["find_force"]
-        find_virial = label_dict["find_virial"]
-        find_atom_ener = label_dict["find_atom_ener"]
+        find_energy = global_cvt_2_ener_float(label_dict["find_energy"])
+        find_force = global_cvt_2_tf_float(label_dict["find_force"])
+        find_virial = global_cvt_2_tf_float(label_dict["find_virial"])
+        find_atom_ener = global_cvt_2_tf_float(label_dict["find_atom_ener"])
 
         if self.enable_atom_ener_coeff:
             # when ener_coeff (\nu) is defined, the energy is defined as
@@ -932,8 +932,8 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
 
         energy_hat = label_dict["energy"]
         ener_dipole_hat = label_dict["energy_dipole"]
-        find_energy = label_dict["find_energy"]
-        find_ener_dipole = label_dict["find_energy_dipole"]
+        find_energy = global_cvt_2_ener_float(label_dict["find_energy"])
+        find_ener_dipole = global_cvt_2_ener_float(label_dict["find_energy_dipole"])
 
         l2_ener_loss = tf.reduce_mean(
             tf.square(energy - energy_hat), name="l2_" + suffix