Skip to content

Commit 4bc6426

Browse files
author
Han Wang
committed
fix(pt): guard against double-loading libdeepmd_op_pt.so
When JAX backend loads libdeepmd_op_pt.so first, a subsequent import deepmd.pt triggers cxx_op.py to load it again, causing a fatal c10::Error double-registration crash. Add a guard that checks if ops are already registered before calling torch.ops.load_library.
1 parent 443a140 commit 4bc6426

1 file changed

Lines changed: 4 additions & 0 deletions

File tree

deepmd/pt/cxx_op.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,10 @@ def load_library(module_name: str) -> bool:
4242
module_file = (SHARED_LIB_DIR / (prefix + module_name)).with_suffix(ext).resolve()
4343

4444
if module_file.is_file():
45+
# Guard: skip if ops are already registered (e.g. by JAX backend
46+
# loading the same .so via a different pathway in the same process).
47+
if hasattr(torch.ops, "deepmd") and hasattr(torch.ops.deepmd, "border_op"):
48+
return True
4549
try:
4650
torch.ops.load_library(module_file)
4751
except OSError as e:

0 commit comments

Comments
 (0)