[BUG] pt: polar training error: KeyError: 'type'

[njzjz]

Bug summary

See below.

DeePMD-kit Version

ed831c8

TensorFlow Version

PT v2.2.0+cu121-g8ac9b20d4b0

How did you download the software?

Built from source

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 807, 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 304, in main
    train(FLAGS)
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/entrypoints/main.py", line 270, in train
    trainer = get_trainer(
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/entrypoints/main.py", line 187, in get_trainer
    trainer = training.Trainer(
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/train/training.py", line 358, in __init__
    single_model_stat(
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/train/training.py", line 248, in single_model_stat
    _model.compute_or_load_stat(
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/model/model/make_model.py", line 482, in compute_or_load_stat
    return self.atomic_model.compute_or_load_stat(sampled_func, stat_file_path)
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/model/atomic_model/dp_atomic_model.py", line 224, in compute_or_load_stat
    self.fitting_net.compute_output_stats(wrapped_sampler, stat_file_path)
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/model/task/polarizability.py", line 233, in compute_output_stats
    nframs = sampled[sys]["type"].shape[0]
KeyError: 'type'

Steps to Reproduce

cd examples/water_tensor/polar
dp --pt train input_torch.json

Further Information, Files, and Links

No response