fix:UTs

deepmodeling · Apr 7, 2024 · 33b33be · 33b33be
1 parent a3ca59e
commit 33b33be
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -199,9 +199,9 @@ def forward_atomic(
                 mapping,
                 fparam,
                 aparam,
-            )["energy"]
+            )
             # apply bias to each individual model
-            ener_list.append(model.apply_out_stat(raw_ret, mapping[extended_atype]))
+            ener_list.append(model.apply_out_stat(raw_ret, mapping[extended_atype][:,:nloc])["energy"])
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)
 
         fit_ret = {

diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
@@ -172,6 +172,10 @@ def forward_common(
             model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
+        @property
+        def out_bias(self)-> torch.Tensor:
+            return self.atomic_model.out_bias
+
         def change_out_bias(
             self,
             merged,