From 52f89fb5ffe5166f686b9cc101a5d976dc3d1e78 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Wed, 25 Sep 2024 23:40:16 +0800
Subject: [PATCH] fix(pt): fix hess calculation for nopbc system

---
 deepmd/pt/model/descriptor/env_mat.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/deepmd/pt/model/descriptor/env_mat.py b/deepmd/pt/model/descriptor/env_mat.py
index e89e7467d3..dc7142249a 100644
--- a/deepmd/pt/model/descriptor/env_mat.py
+++ b/deepmd/pt/model/descriptor/env_mat.py
@@ -21,10 +21,11 @@ def _make_env_mat(
     nall = coord.shape[1]
     mask = nlist >= 0
     # nlist = nlist * mask  ## this impl will contribute nans in Hessian calculation.
-    nlist = torch.where(mask, nlist, nall - 1)
+    nlist = torch.where(mask, nlist, nall)
     coord_l = coord[:, :natoms].view(bsz, -1, 1, 3)
     index = nlist.view(bsz, -1).unsqueeze(-1).expand(-1, -1, 3)
-    coord_r = torch.gather(coord, 1, index)
+    coord_pad = torch.concat([coord, coord[:, -1:, :] + rcut], dim=1)
+    coord_r = torch.gather(coord_pad, 1, index)
     coord_r = coord_r.view(bsz, natoms, nnei, 3)
     diff = coord_r - coord_l
     length = torch.linalg.norm(diff, dim=-1, keepdim=True)