change source of torch_admp from github to pypi in pyproject.toml; minor code improvement

Author: ChiahsinChu

deepmd/pt/modifier/dipole_charge.py

@@ -62,6 +62,13 @@ def __init__(
         self.sys_charge_map = to_torch_tensor(np.array(sys_charge_map))
         self._model_charge_map = model_charge_map
         self._sys_charge_map = sys_charge_map
+        
+        # Validate that model_charge_map and sel_type have matching lengths
+        if len(model_charge_map) != len(sel_type):
+            raise ValueError(
+                f"model_charge_map length ({len(model_charge_map)}) must match "
+                f"sel_type length ({len(sel_type)})"
+            )
 
         # init ewald recp
         self.ewald_h = ewald_h
@@ -344,8 +351,6 @@ def make_mask(
     sel_type = sel_type.to(torch.long)
     atype = atype.to(torch.long)
 
-    # Create mask using broadcasting
-    mask = torch.zeros_like(atype, dtype=torch.bool)
-    for t in sel_type:
-        mask = mask | (atype == t)
+    # Create mask using broadcasting for JIT compatibility
+    mask = (atype.unsqueeze(-1) == sel_type).any(dim=-1)
     return mask

pyproject.toml

@@ -164,11 +164,11 @@ pin_pytorch_cpu = [
   # macos x86 has been deprecated
   "torch>=2.8,<2.10; platform_machine!='x86_64' or platform_system != 'Darwin'",
   "torch; platform_machine=='x86_64' and platform_system == 'Darwin'",
-  "torch_admp @ git+https://github.com/chiahsinchu/torch-admp.git@v1.1.1",
+  "torch-admp==1.1.1",
 ]
 pin_pytorch_gpu = [
   "torch>=2.7,<2.10",
-  "torch_admp @ git+https://github.com/chiahsinchu/torch-admp.git@v1.1.1",
+  "torch-admp==1.1.1",
 ]
 pin_jax = [
   "jax==0.5.0;python_version>='3.10'",