I am not seeing a problem when I do a basic example
train_dataset = []
for _ in range(1000):
    train_dataset.append(torch.rand(4, 4))

def basic_sampler(rank, world_size):
    sampler = torch.utils.data.distributed.DistributedSampler(
    	train_dataset,
    	num_replicas=world_size,
    	rank=rank
    )
    data_loader = DataLoader(train_dataset, batch_size=100, sampler=sampler)
    print(f"rank = {rank} train_dataset sample count = {len(train_dataset)}")
    print(f"rank = {rank} data_loader batch count = {len(list(data_loader))}")


if __name__ == "__main__":
    mp.spawn(basic_sampler, nprocs=2, args=(2,)) 

rank = 0 train_dataset sample count = 1000
rank = 0 data_loader batch count = 5
rank = 1 train_dataset sample count = 1000
rank = 1 data_loader batch count = 5

Are you able to add some debugging logs to confirm that everything is set correctly?