fix bug in model frozen

Author: ChiahsinChu

deepmd/pt/model/model/make_model.py

@@ -297,11 +297,11 @@ def forward_common_lower(
                 comm_dict=comm_dict,
             )
             # add weight to atomic_output
-            kw = next(iter(self.atomic_output_def().var_defs.keys()))
+            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.reshape(
-                    *atomic_ret[kw].shape[:-1], -1
+                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,