reshape

iProzd · Nov 21, 2024 · 47a160f · 47a160f
1 parent b1ae8bc
commit 47a160f
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
@@ -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