DDP with sampler

Hello, when using DDP with the [sampler](https://github.com/microsoft/protein-sequence-models/blob/main/sequence_models/samplers.py) ,  due to the dynamic batch size, the number of batches allocated by different ranks is inconsistent. This is what makes different ranks unable to communicate in DDP. Are there other solutions？

