Skip to content

Commit cca4eb9

Browse files
author
Han Wang
committed
refactor: defuse dpa2 register_buffer trap for tebd compression slot
deepmd/pt_expt/descriptor/dpa2.py:_store_type_embd_data still called torch.nn.Module.register_buffer(self, "type_embd_data", ...) directly, the same raw pattern fixed for dpa1/se_t_tebd in the previous commit. DescrptDPA2 (dpmodel) doesn't declare the type_embd_data slot yet, so this was purely defensive, but it removes the landmine for whichever future change declares it there. Swap to the existing register_buffer_replacing_slot() helper in pt_expt/common.py.
1 parent 1072bb2 commit cca4eb9

1 file changed

Lines changed: 2 additions & 1 deletion

File tree

deepmd/pt_expt/descriptor/dpa2.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
remap_atype_to_padding,
2222
)
2323
from deepmd.pt_expt.common import (
24+
register_buffer_replacing_slot,
2425
torch_module,
2526
)
2627
from deepmd.pt_expt.descriptor.base_descriptor import (
@@ -284,7 +285,7 @@ def _store_type_embd_data(self) -> None:
284285
self.repinit.embeddings_strip[0].call(two_side_embd).detach()
285286
)
286287

287-
torch.nn.Module.register_buffer(self, "type_embd_data", embd_tensor)
288+
register_buffer_replacing_slot(self, "type_embd_data", embd_tensor)
288289

289290
@cast_precision
290291
def call(

0 commit comments

Comments
 (0)