From f261151fbb0e1a21599d69905b393e9ddfbe66d4 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Mon, 23 Mar 2026 12:26:30 +0800
Subject: [PATCH 1/2] fix: Square atom_norm in non-Huber energy and virial loss
 calculations.

---
 deepmd/dpmodel/loss/ener.py | 4 ++--
 deepmd/pd/loss/ener.py      | 4 ++--
 deepmd/pt/loss/ener.py      | 4 ++--
 deepmd/tf/loss/ener.py      | 4 ++--
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py
index 201a77fcdd..8172302f5d 100644
--- a/deepmd/dpmodel/loss/ener.py
+++ b/deepmd/dpmodel/loss/ener.py
@@ -245,7 +245,7 @@ def call(
             if self.loss_func == "mse":
                 l2_ener_loss = xp.mean(xp.square(energy - energy_hat))
                 if not self.use_huber:
-                    loss += atom_norm_ener * (pref_e * l2_ener_loss)
+                    loss += atom_norm_ener**2 * (pref_e * l2_ener_loss)
                 else:
                     l_huber_loss = custom_huber_loss(
                         atom_norm_ener * energy,
@@ -312,7 +312,7 @@ def call(
                     xp.square(virial_hat_reshape - virial_reshape),
                 )
                 if not self.use_huber:
-                    loss += atom_norm * (pref_v * l2_virial_loss)
+                    loss += atom_norm**2 * (pref_v * l2_virial_loss)
                 else:
                     l_huber_loss = custom_huber_loss(
                         atom_norm * virial_reshape,
diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py
index cf093b90d4..9694cd1236 100644
--- a/deepmd/pd/loss/ener.py
+++ b/deepmd/pd/loss/ener.py
@@ -233,7 +233,7 @@ def forward(
                         l2_ener_loss.detach(), find_energy
                     )
                 if not self.use_huber:
-                    loss += atom_norm * (pref_e * l2_ener_loss)
+                    loss += atom_norm**2 * (pref_e * l2_ener_loss)
                 else:
                     l_huber_loss = custom_huber_loss(
                         atom_norm * energy_pred,
@@ -404,7 +404,7 @@ def forward(
                     l2_virial_loss.detach(), find_virial
                 )
             if not self.use_huber:
-                loss += atom_norm * (pref_v * l2_virial_loss)
+                loss += atom_norm**2 * (pref_v * l2_virial_loss)
             else:
                 l_huber_loss = custom_huber_loss(
                     atom_norm * model_pred["virial"].reshape([-1]),
diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py
index 66d60aacec..1830a6e64b 100644
--- a/deepmd/pt/loss/ener.py
+++ b/deepmd/pt/loss/ener.py
@@ -239,7 +239,7 @@ def forward(
                         l2_ener_loss.detach(), find_energy
                     )
                 if not self.use_huber:
-                    loss += atom_norm * (pref_e * l2_ener_loss)
+                    loss += atom_norm**2 * (pref_e * l2_ener_loss)
                 else:
                     l_huber_loss = custom_huber_loss(
                         atom_norm * energy_pred,
@@ -421,7 +421,7 @@ def forward(
                         l2_virial_loss.detach(), find_virial
                     )
                 if not self.use_huber:
-                    loss += atom_norm * (pref_v * l2_virial_loss)
+                    loss += atom_norm**2 * (pref_v * l2_virial_loss)
                 else:
                     l_huber_loss = custom_huber_loss(
                         atom_norm * model_pred["virial"].reshape(-1),
diff --git a/deepmd/tf/loss/ener.py b/deepmd/tf/loss/ener.py
index 91607245a2..726c903dba 100644
--- a/deepmd/tf/loss/ener.py
+++ b/deepmd/tf/loss/ener.py
@@ -346,7 +346,7 @@ def build(
         more_loss = {}
         if self.has_e:
             if not self.use_huber:
-                loss += atom_norm_ener * (pref_e * l2_ener_loss)
+                loss += atom_norm_ener**2 * (pref_e * l2_ener_loss)
             else:
                 l_huber_loss = custom_huber_loss(
                     atom_norm_ener * energy,
@@ -370,7 +370,7 @@ def build(
             )
         if self.has_v:
             if not self.use_huber:
-                loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss))
+                loss += global_cvt_2_ener_float(atom_norm**2 * (pref_v * l2_virial_loss))
             else:
                 l_huber_loss = custom_huber_loss(
                     atom_norm * tf.reshape(virial, [-1]),

From 247f05388fed9b64d49bc940a067780ed1a2b1b8 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 23 Mar 2026 04:28:19 +0000
Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/tf/loss/ener.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/deepmd/tf/loss/ener.py b/deepmd/tf/loss/ener.py
index 726c903dba..08d89fdefd 100644
--- a/deepmd/tf/loss/ener.py
+++ b/deepmd/tf/loss/ener.py
@@ -370,7 +370,9 @@ def build(
             )
         if self.has_v:
             if not self.use_huber:
-                loss += global_cvt_2_ener_float(atom_norm**2 * (pref_v * l2_virial_loss))
+                loss += global_cvt_2_ener_float(
+                    atom_norm**2 * (pref_v * l2_virial_loss)
+                )
             else:
                 l_huber_loss = custom_huber_loss(
                     atom_norm * tf.reshape(virial, [-1]),