fix reshape

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
iProzd · Nov 21, 2024 · b1ae8bc · b1ae8bc
1 parent 6ec5f6f
commit b1ae8bc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd/dpmodel/descriptor/se_t_tebd.py b/deepmd/dpmodel/descriptor/se_t_tebd.py
@@ -737,7 +737,7 @@ def call(
             assert type_embedding is not None
             ntypes_with_padding = type_embedding.shape[0]
             # nf x (nl x nnei)
-            nlist_index = nlist_masked.reshape(nf, nloc * nnei)
+            nlist_index = xp.reshape(nlist_masked, (nf, nloc * nnei))
             # nf x (nl x nnei)
             nei_type = xp_take_along_axis(atype_ext, nlist_index, axis=1)
             # nfnl x nnei