[BUG] pt se_e2_r failed to do env mat stat

[njzjz]

Bug summary

When training a se_e2_r model, get the error.

DeePMD-kit Version

latest

TensorFlow Version

PyTorch v2.2.0+cu121-g8ac9b20d4b0

How did you download the software?

Offline packages

Input Files, Running Commands, Error Log, etc.

Traceback (most recent call last):
  File "/home/jz748/anaconda3/bin/dp", line 8, in <module>
    sys.exit(main())
  File "/home/jz748/codes/deepmd-kit/deepmd/main.py", line 805, in main
    deepmd_main(args)
  File "/home/jz748/anaconda3/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 347, in wrapper
    return f(*args, **kwargs)
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/entrypoints/main.py", line 282, in main
    train(FLAGS)
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/entrypoints/main.py", line 246, in train
    trainer = get_trainer(
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/entrypoints/main.py", line 193, in get_trainer
    trainer = training.Trainer(
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/train/training.py", line 242, in __init__
    self.model = get_single_model(model_params, sampled, stat_file_path)
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/train/training.py", line 189, in get_single_model
    model.compute_or_load_stat(
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/model/atomic_model/dp_atomic_model.py", line 194, in compute_or_load_stat
    self.descriptor.compute_input_stats(sampled, stat_file_path)
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/model/descriptor/se_r.py", line 154, in compute_input_stats
    env_mat_stat.load_or_compute_stats(merged, path)
  File "/home/jz748/codes/deepmd-kit/deepmd/utils/env_mat_stat.py", line 174, in load_or_compute_stats
    self.compute_stats(data)
  File "/home/jz748/codes/deepmd-kit/deepmd/utils/env_mat_stat.py", line 108, in compute_stats
    for iter_stats in self.iter(data):
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/utils/env_mat_stat.py", line 136, in iter
    env_mat = env_mat.view(
RuntimeError: shape '[1920, 138, 1]' is invalid for input of size 1059840
terminate called after throwing an instance of 'std::system_error'
  what():  Broken pipe
terminate called after throwing an instance of 'std::system_error'
  what():  Broken pipe
terminate called after throwing an instance of 'std::system_error'
  what():  Broken pipe
terminate called after throwing an instance of 'std::system_error'
  what():  Broken pipe
terminate called after throwing an instance of 'std::system_error'
  what():  Broken pipe

Steps to Reproduce

cd examples/water/se_e2_r
dp --pt train input.json

Further Information, Files, and Links

No response

[anyangml]

Resolved in #3351.