chore:use forward_common_atomic

deepmodeling · Apr 7, 2024 · 90b8d56 · 90b8d56
1 parent 89fe042
commit 90b8d56
Showing 1 changed file with 4 additions and 7 deletions.
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -190,19 +190,16 @@ def forward_atomic(
 
         for i, model in enumerate(self.models):
             mapping = self.mapping_list[i]
-            raw_ret = model.forward_atomic(
+            # apply bias to each individual model
+            ener_list.append(
+                model.forward_common_atomic(
                 extended_coord,
                 mapping[extended_atype],
                 nlists_[i],
                 mapping,
                 fparam,
                 aparam,
-            )
-            # apply bias to each individual model
-            ener_list.append(
-                model.apply_out_stat(raw_ret, mapping[extended_atype][:, :nloc])[
-                    "energy"
-                ]
+            )["energy"]
             )
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)