fix: move _load_custom_ops after deepmd.pt import in gen scripts

The _load_custom_ops() call was placed before deepmd.pt was imported,
so the guard `hasattr(torch.ops.deepmd, "border_op")` didn't see the
op registered by the standard install path. This caused the build
directory's .so to be loaded, and then when deepmd.pt was imported
later for .pth export, it loaded the same op again — crash.

Fix: move _load_custom_ops() to after the deepmd.pt import so the
guard correctly detects the already-registered op and skips.

Author: Han Wang

source/tests/infer/gen_dpa1.py

@@ -100,8 +100,6 @@ def main():
         get_model,
     )
 
-    # Load custom ops after deepmd import to avoid double registration
-    _load_custom_ops()
     _ensure_inductor_compiler()
 
     # ---- 1. DPA1 model config with type_one_side=True ----
@@ -152,6 +150,9 @@ def main():
         deserialize_to_file as pt_expt_deserialize_to_file,
     )
 
+    # Load custom ops after deepmd.pt import to avoid double registration
+    _load_custom_ops()
+
     base_dir = os.path.dirname(__file__)
 
     pt2_path = os.path.join(base_dir, "deeppot_dpa1.pt2")

source/tests/infer/gen_dpa2.py

@@ -100,8 +100,6 @@ def main():
         get_model,
     )
 
-    # Load custom ops after deepmd import to avoid double registration
-    _load_custom_ops()
     _ensure_inductor_compiler()
 
     # ---- 1. DPA2 model config with type_one_side=True, use_three_body=True ----
@@ -175,6 +173,9 @@ def main():
         deserialize_to_file as pt_expt_deserialize_to_file,
     )
 
+    # Load custom ops after deepmd.pt import to avoid double registration
+    _load_custom_ops()
+
     base_dir = os.path.dirname(__file__)
 
     pt2_path = os.path.join(base_dir, "deeppot_dpa2.pt2")

source/tests/infer/gen_dpa3.py

@@ -105,8 +105,6 @@ def main():
         get_model,
     )
 
-    # Load custom ops after deepmd import to avoid double registration
-    _load_custom_ops()
     _ensure_inductor_compiler()
 
     # ---- 1. DPA3 model config (small, fast to compile) ----
@@ -158,6 +156,9 @@ def main():
         deserialize_to_file as pt_expt_deserialize_to_file,
     )
 
+    # Load custom ops after deepmd.pt import to avoid double registration
+    _load_custom_ops()
+
     base_dir = os.path.dirname(__file__)
 
     pt2_path = os.path.join(base_dir, "deeppot_dpa3.pt2")

source/tests/infer/gen_fparam_aparam.py

@@ -75,8 +75,6 @@ def main():
         get_model,
     )
 
-    # Load custom ops after deepmd import to avoid double registration
-    _load_custom_ops()
     _ensure_inductor_compiler()
 
     # ---- 1. Model config (type_one_side=True) ----
@@ -134,6 +132,9 @@ def main():
         deserialize_to_file as pt_expt_deserialize_to_file,
     )
 
+    # Load custom ops after deepmd.pt import to avoid double registration
+    _load_custom_ops()
+
     base_dir = os.path.dirname(__file__)
 
     pt2_path = os.path.join(base_dir, "fparam_aparam.pt2")

source/tests/infer/gen_sea.py

@@ -78,7 +78,7 @@ def main():
         deserialize_to_file,
     )
 
-    # Load custom ops after deepmd import to avoid double registration
+    # Load custom ops after deepmd.pt import to avoid double registration
     _load_custom_ops()
     _ensure_inductor_compiler()