Perf: replace unnecessary torch.split with indexing (#4505)

Some operations only use the first segment of the result tensor of
`torch.split`. In this case, all the other segments are created and
discarded. This slightly adds an overhead to the training process.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **Bug Fixes**
- Simplified tensor slicing operations in the `RepformerLayer` class and
the `nlist_distinguish_types` function, enhancing readability and
performance.
  
- **Documentation**
- Updated comments for clarity regarding tensor shapes in the
`RepformerLayer` class.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: caic99

deepmd/pt/model/descriptor/repformer_layer.py

@@ -1003,7 +1003,7 @@ def _cal_grrg(h2g2: torch.Tensor, axis_neuron: int) -> torch.Tensor:
         # nb x nloc x 3 x ng2
         nb, nloc, _, ng2 = h2g2.shape
         # nb x nloc x 3 x axis
-        h2g2m = torch.split(h2g2, axis_neuron, dim=-1)[0]
+        h2g2m = h2g2[..., :axis_neuron]
         # nb x nloc x axis x ng2
         g1_13 = torch.matmul(torch.transpose(h2g2m, -1, -2), h2g2) / (3.0**1)
         # nb x nloc x (axisxng2)

deepmd/pt/utils/nlist.py

@@ -310,7 +310,7 @@ def nlist_distinguish_types(
         inlist = torch.gather(nlist, 2, imap)
         inlist = inlist.masked_fill(~(pick_mask.to(torch.bool)), -1)
         # nloc x nsel[ii]
-        ret_nlist.append(torch.split(inlist, [ss, snsel - ss], dim=-1)[0])
+        ret_nlist.append(inlist[..., :ss])
     return torch.concat(ret_nlist, dim=-1)