fix code for view ops

Author: HydrogenSulfate

deepmd/pd/model/descriptor/repflows.py

@@ -604,7 +604,10 @@ def forward(
                     _shapes = node_ebd.shape[1:]
                     _shapes[1] = n_padding
                     node_ebd = paddle.concat(
-                        [node_ebd, paddle.zeros(_shapes, dtype=node_ebd.dtype)],
+                        [
+                            node_ebd.squeeze(0),
+                            paddle.zeros(_shapes, dtype=node_ebd.dtype),
+                        ],
                         axis=1,
                     )
                     real_nloc = nloc
@@ -652,7 +655,7 @@ def forward(
                         place=paddle.CPUPlace(),
                     ),  # should be int of c++, placed on cpu
                 )
-                node_ebd_ext = ret[0].unsqueeze(0)
+                node_ebd_ext = paddle.assign(ret).unsqueeze(0)
                 if has_spin:
                     node_ebd_real_ext, node_ebd_virtual_ext = paddle.split(
                         node_ebd_ext, [n_dim, n_dim], axis=2

deepmd/pd/model/descriptor/repformers.py

@@ -536,11 +536,7 @@ def forward(
                         place=paddle.CPUPlace(),
                     ),  # should be int of c++, placed on cpu
                 )
-                # print(f"ret.shape = {ret.shape}")
-                # print(f"ret[0].shape = ", ret[0].shape)
-                g1_ext = ret.unsqueeze(0)
-                # print(f"g1_ext.shape = ", g1_ext.shape)
-                # exit()
+                g1_ext = paddle.assign(ret).unsqueeze(0)
                 if has_spin:
                     g1_real_ext, g1_virtual_ext = paddle.split(
                         g1_ext, [ng1, ng1], dim=2