add doc for atomic_weight

Author: ChiahsinChu

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -171,6 +171,9 @@ def forward_common_atomic(
             frame parameters, shape: nf x dim_fparam
         aparam
             atomic parameter, shape: nf x nloc x dim_aparam
+        atomic_weight
+            atomic weights for scaling outputs, shape: nf x nloc x dim_aw
+            if provided, all output values will be multiplied by this weight.
 
         Returns
         -------

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -227,6 +227,9 @@ def forward_common_atomic(
             atomic parameter, shape: nf x nloc x dim_aparam
         comm_dict
             The data needed for communication for parallel inference.
+        atomic_weight
+            atomic weights for scaling outputs, shape: nf x nloc x dim_aw
+            if provided, all output values will be multiplied by this weight.
 
         Returns
         -------