[BUG] Incomplete and risky bias statistics

[iProzd]

Bug summary

In data statistics, certain types may not be sampled from the dataset, resulting in incomplete bias statistics. This will cause training problems, especially when dealing with mixed-type data formats.

The PyTorch DataLoader could be enhanced by implementing two methods:

1. Calculate the number of atoms of each type within the dataset and cache some frame indices for each type.
2. If the sampled frames lack certain types that exist in the dataset, use the cached indices to add frames of these missing types into the samples before performing bias statistics.

This approach will ensure comprehensive bias statistics.

DeePMD-kit Version

3.0.0

Backend and its version

Both Pytorch and TensorFlow

How did you download the software?

Built from source

Input Files, Running Commands, Error Log, etc.

See above

Steps to Reproduce

See above

Further Information, Files, and Links

No response