From 1de75cd1856b34842583189a7c208047a42e6fff Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Wed, 1 Jul 2026 15:08:16 +0800
Subject: [PATCH] fix(tf): preserve loss config across linear submodels

LinearModel.get_loss reassigned its `loss` argument to each submodel's
return value. Frozen and pairtab submodels return None from get_loss, so when
such a non-trainable submodel preceded a trainable energy submodel, the next
submodel was called with loss=None and EnerFitting.get_loss crashed on
None.pop("type").

Iterate with a separate variable and pass a deepcopy of the original loss
config to each submodel, returning the first non-None result. Add a test with
a None-returning submodel ordered before a trainable one.

Fix #5682
---
 deepmd/tf/model/linear.py            | 12 ++++++----
 source/tests/tf/test_linear_model.py | 35 ++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/deepmd/tf/model/linear.py b/deepmd/tf/model/linear.py
index 90f559f73f..4d147d5169 100644
--- a/deepmd/tf/model/linear.py
+++ b/deepmd/tf/model/linear.py
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
 import operator
 from enum import (
     Enum,
@@ -77,11 +78,14 @@ def get_fitting(self) -> Fitting | dict:
 
     def get_loss(self, loss: dict, lr: LearningRateExp) -> Loss | dict | None:
         """Get the loss function(s)."""
-        # the first model that is not None, or None if all models are None
+        # Return the first submodel loss that is not None, or None if all
+        # submodels are non-trainable. Each submodel must receive the original
+        # loss config: frozen/table submodels return None, so consuming the
+        # return value would pass None to a later trainable submodel.
         for model in self.models:
-            loss = model.get_loss(loss, lr)
-            if loss is not None:
-                return loss
+            submodel_loss = model.get_loss(copy.deepcopy(loss), lr)
+            if submodel_loss is not None:
+                return submodel_loss
         return None
 
     def get_rcut(self) -> float:
diff --git a/source/tests/tf/test_linear_model.py b/source/tests/tf/test_linear_model.py
index 6743c4658b..1392a97820 100644
--- a/source/tests/tf/test_linear_model.py
+++ b/source/tests/tf/test_linear_model.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import os
+import unittest
 
 import numpy as np
 
@@ -121,3 +122,37 @@ def tearDown(self) -> None:
         for pb in self.graph_dirs:
             os.remove(pb)
         del_data()
+
+
+class TestLinearModelGetLoss(unittest.TestCase):
+    """A non-trainable submodel (frozen/pairtab) returns ``None`` from
+    ``get_loss``; a later trainable submodel must still receive the original
+    loss configuration rather than that ``None``.
+    """
+
+    def test_get_loss_skips_none_returning_submodels(self) -> None:
+        class _NoneLossModel:
+            """Mimics frozen/pairtab submodels, which own no trainable loss."""
+
+            def get_loss(self, loss, lr):
+                return None
+
+        class _EnerLossModel:
+            """Mimics a trainable ener submodel: dereferences loss as a dict
+            (as EnerFitting.get_loss does via ``loss.pop``).
+            """
+
+            def get_loss(self, loss, lr):
+                loss.pop("type", "ener")
+                return "ener-loss"
+
+        # bypass the heavy __init__; get_loss only touches self.models
+        model = object.__new__(LinearEnergyModel)
+        model.models = [_NoneLossModel(), _EnerLossModel()]
+
+        loss_config = {"type": "ener", "start_pref_e": 1.0}
+        result = model.get_loss(loss_config, lr=None)
+
+        self.assertEqual(result, "ener-loss")
+        # the original config must be preserved for other consumers
+        self.assertEqual(loss_config, {"type": "ener", "start_pref_e": 1.0})