From 47a160f33614d7c1dbb02431dd9724831b5bd14d Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 21 Nov 2024 15:02:02 -0500 Subject: [PATCH] reshape --- deepmd/dpmodel/descriptor/se_t_tebd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/dpmodel/descriptor/se_t_tebd.py b/deepmd/dpmodel/descriptor/se_t_tebd.py index 530e0cc98b..c350e3eb47 100644 --- a/deepmd/dpmodel/descriptor/se_t_tebd.py +++ b/deepmd/dpmodel/descriptor/se_t_tebd.py @@ -775,7 +775,7 @@ def call( gg_t = xp_take_along_axis(tt_full, idx, axis=0) # (nfnl x nt_i x nt_j) x ng - gg_t = gg_t.reshape(nf * nloc, nnei, nnei, ng) + gg_t = xp.reshape(gg_t, (nf * nloc, nnei, nnei, ng)) if self.smooth: gg_t = ( gg_t