Fix: pt tensor loss label name (#4587)

To address polar nan loss mentioned in #4586 

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **Refactor**
- Adjusted the internal processing order in computation routines to
enhance consistency while maintaining the same overall user experience.
- Updated model prediction handling to ensure compatibility in shape
during statistical computations, reducing potential runtime errors.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: anyangml

deepmd/pt/train/training.py

@@ -1253,13 +1253,11 @@ def get_loss(loss_params, start_lr, _ntypes, _model):
         if "mask" in model_output_type:
             model_output_type.pop(model_output_type.index("mask"))
         tensor_name = model_output_type[0]
-        loss_params["tensor_name"] = tensor_name
         loss_params["tensor_size"] = _model.model_output_def()[tensor_name].output_size
-        label_name = tensor_name
-        if label_name == "polarizability":
-            label_name = "polar"
-        loss_params["label_name"] = label_name
-        loss_params["tensor_name"] = label_name
+        loss_params["label_name"] = tensor_name
+        if tensor_name == "polarizability":
+            tensor_name = "polar"
+        loss_params["tensor_name"] = tensor_name
         return TensorLoss(**loss_params)
     elif loss_type == "property":
         task_dim = _model.get_task_dim()

deepmd/pt/utils/stat.py

@@ -477,7 +477,9 @@ def compute_output_stats_global(
         # subtract the model bias and output the delta bias
 
         stats_input = {
-            kk: merged_output[kk] - model_pred[kk] for kk in keys if kk in merged_output
+            kk: merged_output[kk] - model_pred[kk].reshape(merged_output[kk].shape)
+            for kk in keys
+            if kk in merged_output
         }
 
     bias_atom_e = {}