Using customized ddp sampler get stuck during training

[GongXinyuu]

I have set trainer. replace_sampler_ddp=False and implement a very naive DistributedSampler as below:

def create_normal_sampler(dataset, shuffle, drop_last) -> Optional[Sampler]:
    if dist.is_initialized():
        logger.info("DistributedSampler is enabled.")
        return torch.utils.data.distributed.DistributedSampler(
            dataset, shuffle=shuffle, drop_last=drop_last
        )
    else:
        return None

sampler = create_normal_sampler(
    dataset, split_cfg.shuffle, split_cfg.drop_last
)
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=split_cfg.batch_size,
    shuffle=split_cfg.shuffle if sampler is None else False,
    num_workers=split_cfg.num_workers,
    drop_last=split_cfg.drop_last,
    pin_memory=split_cfg.pin_memory,
    collate_fn=collate_fn,
    sampler=sampler,
)

The training gets stuck without any error message. Before I do that, everything runs well.