fix 2nd grad is NaN for torch.norm at 0

Author: caic99

deepmd/pt/model/descriptor/repflows.py

@@ -465,8 +465,9 @@ def forward(
         # beyond the cutoff sw should be 0.0
         sw = sw.masked_fill(~nlist_mask, 0.0)
 
+        eps = 1e-6  # for numerical stability; see #4809
         # get angle nlist (maybe smaller)
-        a_dist_mask = (torch.linalg.norm(diff, dim=-1) < self.a_rcut)[
+        a_dist_mask = (torch.linalg.norm(diff + eps, dim=-1) < self.a_rcut)[
             :, :, : self.a_sel
         ]
         a_nlist = nlist[:, :, : self.a_sel]
@@ -505,17 +506,16 @@ def forward(
         edge_input, h2 = torch.split(dmatrix, [1, 3], dim=-1)
         if self.edge_init_use_dist:
             # nb x nloc x nnei x 1
-            edge_input = torch.linalg.norm(diff, dim=-1, keepdim=True)
+            edge_input = torch.linalg.norm(diff + eps, dim=-1, keepdim=True)
 
         # nf x nloc x a_nnei x 3
         normalized_diff_i = a_diff / (
-            torch.linalg.norm(a_diff, dim=-1, keepdim=True) + 1e-6
+            torch.linalg.norm(a_diff + eps, dim=-1, keepdim=True) + eps
         )
         # nf x nloc x 3 x a_nnei
         normalized_diff_j = torch.transpose(normalized_diff_i, 2, 3)
         # nf x nloc x a_nnei x a_nnei
-        # 1 - 1e-6 for torch.acos stability
-        cosine_ij = torch.matmul(normalized_diff_i, normalized_diff_j) * (1 - 1e-6)
+        cosine_ij = torch.matmul(normalized_diff_i, normalized_diff_j)
         angle_input = cosine_ij.unsqueeze(-1) / (torch.pi**0.5)
 
         if not parallel_mode and self.use_loc_mapping: