clear

Author: OutisLi

deepmd/tf/entrypoints/change_bias.py

@@ -199,6 +199,12 @@ def _change_bias_checkpoint_file(
     if stop_batch is None and num_epoch is not None:
         if num_epoch <= 0:
             raise ValueError("training.num_epoch must be positive.")
+        # Apply sys_probs and auto_prob from original training config
+        # to ensure stop_batch calculation matches the original training
+        training_data_config = training_params.get("training_data", {})
+        sys_probs = training_data_config.get("sys_probs", None)
+        auto_prob = training_data_config.get("auto_prob", "prob_sys_size")
+        data.set_sys_probs(sys_probs=sys_probs, auto_prob_style=auto_prob)
         total_numb_batch = compute_total_numb_batch(data.nbatches, data.sys_probs)
         if total_numb_batch <= 0:
             raise ValueError("Total number of training batches must be positive.")

deepmd/utils/argcheck.py

@@ -3672,11 +3672,20 @@ def training_extra_check(data: dict | None) -> bool:
                     raise ValueError(
                         "training.num_epoch_dict is mutually exclusive with training.model_prob."
                     )
+            else:
+                if num_steps is None:
+                    raise ValueError(
+                        "Multi-task mode requires either training.numb_steps or training.num_epoch_dict."
+                    )
         else:
             if num_steps is not None and num_epoch is not None:
                 raise ValueError(
                     "training.num_step and training.num_epoch are mutually exclusive."
                 )
+            if num_steps is None and num_epoch is None:
+                raise ValueError(
+                    "Single-task mode requires either training.numb_steps or training.num_epoch."
+                )
         return True
 
     doc_training = "The training options."

source/tests/pt/test_sampler.py

@@ -440,8 +440,6 @@ def test_num_epoch_dict(self) -> None:
         sampler_2 = pt_dataloader.get_sampler_from_params(
             dataset_2, {"sys_probs": [0.4, 0.6], "auto_prob": "prob_sys_size"}
         )
-        probs_1 = self._normalize_probs(np.asarray(sampler_1.weights))
-        probs_2 = self._normalize_probs(np.asarray(sampler_2.weights))
 
         # === Step 2. Compute per-task total_numb_batch ===
         per_task_total = np.array(