From db4ea9656d0ac556da933b7ca16c836827caab18 Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model: gpt-5.2))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Sun, 15 Mar 2026 16:43:03 +0000
Subject: [PATCH 1/2] fix(dpmodel): align Loss.call return type with
 implementation

Authored by OpenClaw (model: gpt-5.2)
---
 deepmd/dpmodel/loss/ener.py |  2 +-
 deepmd/dpmodel/loss/loss.py | 12 ++++++++++--
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py
index 9ab141bdc2..7c9f82588b 100644
--- a/deepmd/dpmodel/loss/ener.py
+++ b/deepmd/dpmodel/loss/ener.py
@@ -93,7 +93,7 @@ def call(
         natoms: int,
         model_dict: dict[str, Array],
         label_dict: dict[str, Array],
-    ) -> dict[str, Array]:
+    ) -> tuple[Array, dict[str, Array]]:
         """Calculate loss from model results and labeled results."""
         energy = model_dict["energy"]
         force = model_dict["force"]
diff --git a/deepmd/dpmodel/loss/loss.py b/deepmd/dpmodel/loss/loss.py
index 4b9831c344..1037c1e469 100644
--- a/deepmd/dpmodel/loss/loss.py
+++ b/deepmd/dpmodel/loss/loss.py
@@ -28,8 +28,16 @@ def call(
         natoms: int,
         model_dict: dict[str, Array],
         label_dict: dict[str, Array],
-    ) -> dict[str, Array]:
-        """Calculate loss from model results and labeled results."""
+    ) -> tuple[Array, dict[str, Array]]:
+        """Calculate loss from model results and labeled results.
+
+        Returns
+        -------
+        loss
+            The scalar loss to minimize.
+        more_loss
+            A dictionary of additional loss terms/metrics for logging.
+        """
 
     @property
     @abstractmethod

From 5979d62cac71bbb9cd18bc0d725ea0116bad6253 Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model: gpt-5.2))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Sun, 15 Mar 2026 17:02:00 +0000
Subject: [PATCH 2/2] fix(dpmodel): avoid natoms[0] in generalized-force loss

Authored by OpenClaw (model: gpt-5.2)
---
 deepmd/dpmodel/loss/ener.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py
index 7c9f82588b..e0d9b39f57 100644
--- a/deepmd/dpmodel/loss/ener.py
+++ b/deepmd/dpmodel/loss/ener.py
@@ -244,10 +244,10 @@ def call(
         if self.has_gf:
             find_drdq = label_dict["find_drdq"]
             drdq = label_dict["drdq"]
-            force_reshape_nframes = xp.reshape(force, (-1, natoms[0] * 3))
-            force_hat_reshape_nframes = xp.reshape(force_hat, (-1, natoms[0] * 3))
+            force_reshape_nframes = xp.reshape(force, (-1, natoms * 3))
+            force_hat_reshape_nframes = xp.reshape(force_hat, (-1, natoms * 3))
             drdq_reshape = xp.reshape(
-                drdq, (-1, natoms[0] * 3, self.numb_generalized_coord)
+                drdq, (-1, natoms * 3, self.numb_generalized_coord)
             )
             gen_force_hat = xp.einsum(
                 "bij,bi->bj", drdq_reshape, force_hat_reshape_nframes