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