move atomic_weight to atomic model level

Author: ChiahsinChu

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -203,6 +203,7 @@ def forward_common_atomic(
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         comm_dict: Optional[dict[str, torch.Tensor]] = None,
+        atomic_weight: Optional[torch.Tensor] = None,
     ) -> dict[str, torch.Tensor]:
         """Common interface for atomic inference.
 
@@ -271,6 +272,10 @@ def forward_common_atomic(
                 ret_dict[kk].reshape([out_shape[0], out_shape[1], out_shape2])
                 * atom_mask[:, :, None]
             ).view(out_shape)
+            if atomic_weight is not None:
+                ret_dict[kk] = ret_dict[kk] * atomic_weight.view(
+                    [out_shape[0], out_shape[1], -1]
+                )
         ret_dict["mask"] = atom_mask
 
         return ret_dict
@@ -284,6 +289,7 @@ def forward(
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         comm_dict: Optional[dict[str, torch.Tensor]] = None,
+        atomic_weight: Optional[torch.Tensor] = None,
     ) -> dict[str, torch.Tensor]:
         return self.forward_common_atomic(
             extended_coord,
@@ -293,6 +299,7 @@ def forward(
             fparam=fparam,
             aparam=aparam,
             comm_dict=comm_dict,
+            atomic_weight=atomic_weight,
         )
 
     def change_type_map(

deepmd/pt/model/model/make_model.py

@@ -295,16 +295,8 @@ def forward_common_lower(
                 fparam=fp,
                 aparam=ap,
                 comm_dict=comm_dict,
+                atomic_weight=atomic_weight,
             )
-            # add weight to atomic_output
-            if hasattr(self.atomic_model, "fitting_net"):
-                if hasattr(self.atomic_model.fitting_net, "var_name"):
-                    kw = self.atomic_model.fitting_net.var_name
-                    if atomic_weight is not None:
-                        # atomic_weight: nf x nloc x dim
-                        atomic_ret[kw] = atomic_ret[kw] * atomic_weight.view(
-                            [atomic_ret[kw].shape[0], atomic_ret[kw].shape[1], -1]
-                        )
             model_predict = fit_output_to_model_output(
                 atomic_ret,
                 self.atomic_output_def(),