fix(pt): pairtab (#5119)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Improved numerical stability of pairwise distance computations to
prevent NaN gradients for zero differences and padded/masked entries;
outputs remain unchanged while gradient behavior is now robust for those
edge cases.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: OutisLi <137472077+OutisLi@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: OutisLi

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -392,14 +392,29 @@ def _get_pairwise_dist(coords: torch.Tensor, nlist: torch.Tensor) -> torch.Tenso
         -------
         torch.Tensor
             The pairwise distance between the atoms (nframes, nloc, nnei).
+
+        Notes
+        -----
+        Safe gradient implementation: when diff is zero (padding entries),
+            both distance and gradient are zero.
         """
         nframes, nloc, nnei = nlist.shape
         coord_l = coords[:, :nloc].view(nframes, -1, 1, 3)
         index = nlist.view(nframes, -1).unsqueeze(-1).expand(-1, -1, 3)
         coord_r = torch.gather(coords, 1, index)
         coord_r = coord_r.view(nframes, nloc, nnei, 3)
         diff = coord_r - coord_l
-        pairwise_rr = torch.linalg.norm(diff, dim=-1, keepdim=True).squeeze(-1)
+        diff_sq = torch.sum(diff * diff, dim=-1, keepdim=True)
+
+        # When diff is zero, output is zero and gradient is also zero
+        mask = diff_sq.squeeze(-1) > 0
+        pairwise_rr = torch.where(
+            mask.unsqueeze(-1),
+            torch.sqrt(
+                torch.where(mask.unsqueeze(-1), diff_sq, torch.ones_like(diff_sq))
+            ),
+            torch.zeros_like(diff_sq),
+        ).squeeze(-1)
         return pairwise_rr
 
     @staticmethod