fix reshape

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
deepmodeling · Sep 25, 2024 · e7aeca0 · e7aeca0
1 parent aff0b42
commit e7aeca0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepmd/dpmodel/utils/exclude_mask.py b/deepmd/dpmodel/utils/exclude_mask.py
@@ -127,7 +127,8 @@ def build_type_exclude_mask(
         index = xp.reshape(
             xp.where(nlist == -1, xp.full_like(nlist, nall), nlist), (nf, nloc * nnei)
         )
-        type_j = xp_take_along_axis(ae, index, axis=1).reshape(nf, nloc, nnei)
+        type_j = xp_take_along_axis(ae, index, axis=1)
+        type_j = xp.reshape(type_j, (nf, nloc, nnei))
         type_ij = type_i[:, :, None] + type_j
         # nf x (nloc x nnei)
         type_ij = xp.reshape(type_ij, (nf, nloc * nnei))