fix the squeeze issue for atomic virial, pt_expt backend only

Author: Han Wang

deepmd/pt_expt/model/ener_model.py

@@ -58,7 +58,7 @@ def forward(
         if self.do_grad_c("energy"):
             model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
             if do_atomic_virial:
-                model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)
+                model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2)
         if "mask" in model_ret:
             model_predict["mask"] = model_ret["mask"]
         return model_predict
@@ -91,7 +91,7 @@ def _forward_lower(
             model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
             if do_atomic_virial:
                 model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze(
-                    -3
+                    -2
                 )
         if "mask" in model_ret:
             model_predict["mask"] = model_ret["mask"]