make exclusion mask modules

Author: Han Wang

deepmd/pt_expt/utils/exclude_mask.py

@@ -12,15 +12,33 @@
 )
 
 
-class AtomExcludeMask(AtomExcludeMaskDP):
+class AtomExcludeMask(AtomExcludeMaskDP, torch.nn.Module):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        torch.nn.Module.__init__(self)
+        AtomExcludeMaskDP.__init__(self, *args, **kwargs)
+
     def __setattr__(self, name: str, value: Any) -> None:
-        if name == "type_mask":
+        if name == "type_mask" and "_buffers" in self.__dict__:
             value = None if value is None else torch.as_tensor(value, device=env.DEVICE)
+            if name in self._buffers:
+                self._buffers[name] = value
+                return
+            self.register_buffer(name, value)
+            return
         return super().__setattr__(name, value)
 
 
-class PairExcludeMask(PairExcludeMaskDP):
+class PairExcludeMask(PairExcludeMaskDP, torch.nn.Module):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        torch.nn.Module.__init__(self)
+        PairExcludeMaskDP.__init__(self, *args, **kwargs)
+
     def __setattr__(self, name: str, value: Any) -> None:
-        if name == "type_mask":
+        if name == "type_mask" and "_buffers" in self.__dict__:
             value = None if value is None else torch.as_tensor(value, device=env.DEVICE)
+            if name in self._buffers:
+                self._buffers[name] = value
+                return
+            self.register_buffer(name, value)
+            return
         return super().__setattr__(name, value)

source/tests/pt_expt/utils/test_exclusion_mask.py

@@ -39,6 +39,10 @@ def test_build_type_exclude_mask(self) -> None:
         mask = des.build_type_exclude_mask(torch.as_tensor(atype, device=env.DEVICE))
         np.testing.assert_equal(mask.detach().cpu().numpy(), expected_mask)
 
+    def test_type_mask_is_buffer(self) -> None:
+        des = AtomExcludeMask(3, exclude_types=[0])
+        assert "type_mask" in des.state_dict()
+
 
 class TestPairExcludeMask(unittest.TestCase, TestCaseSingleFrameWithNlist):
     def setUp(self) -> None:
@@ -62,3 +66,7 @@ def test_build_type_exclude_mask(self) -> None:
             torch.as_tensor(self.atype_ext, device=env.DEVICE),
         )
         np.testing.assert_equal(mask.detach().cpu().numpy(), expected_mask)
+
+    def test_type_mask_is_buffer(self) -> None:
+        des = PairExcludeMask(self.nt, exclude_types=[[0, 1]])
+        assert "type_mask" in des.state_dict()