sakh251:

But for DistributedSampler, I set 3 and rank from 0-3.


I think this might be the issue, for num_replicas=3, the ranks should be from 0-2 and not 0-3. I was able to reproduce the error you’re seeing if I do something like this (which is invalid):
sampler = DistributedSampler(dataset, num_replicas=3, rank=3)