fix the shape of idx_i and idx_j

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
deepmodeling · Nov 21, 2024 · 65da179 · 65da179
1 parent 289ebac
commit 65da179
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py
@@ -965,7 +965,7 @@ def call(
                     ),
                     (-1),
                 )
-                idx_j = nei_type
+                idx_j = xp.reshape(nei_type, (-1,))
                 # (nf x nl x nnei) x ng
                 idx = xp.tile(xp.reshape((idx_i + idx_j), (-1, 1)), (1, ng))
                 # (ntypes) * ntypes * nt

diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py
@@ -569,7 +569,7 @@ def forward(
                 idx_i = torch.tile(
                     atype.reshape(-1, 1) * ntypes_with_padding, [1, nnei]
                 ).view(-1)
-                idx_j = nei_type
+                idx_j = nei_type.view(-1)
                 # (nf x nl x nnei) x ng
                 idx = (
                     (idx_i + idx_j)