fix(pt-expt): clear params on None

Author: Han Wang

deepmd/pt_expt/utils/network.py

@@ -48,6 +48,12 @@ def __setattr__(self, name: str, value: Any) -> None:
         if name in {"w", "b", "idt"} and "_parameters" in self.__dict__:
             val = to_torch_array(value)
             if val is None:
+                if name in self._parameters:
+                    self._parameters[name] = None
+                    return
+                if name in self._buffers:
+                    self._buffers[name] = None
+                    return
                 return super().__setattr__(name, None)
             if getattr(self, "trainable", False):
                 param = (

source/tests/pt_expt/utils/test_network.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+from deepmd.pt_expt.utils.network import (
+    NativeLayer,
+)
+
+
+def test_native_layer_clears_parameter_on_none() -> None:
+    layer = NativeLayer(2, 3, trainable=True)
+    assert layer.w is not None
+    layer.w = None
+    assert layer.w is None
+    assert layer._parameters.get("w") is None
+
+
+def test_native_layer_clears_buffer_on_none() -> None:
+    layer = NativeLayer(2, 3, trainable=False)
+    assert layer.w is not None
+    layer.w = None
+    assert layer.w is None
+    assert layer._buffers.get("w") is None