Training hangs at the end of epoch, with set_transform/with_transform+multiple workers

## Describe the bug
I use load_dataset() (I tried with [wiki](https://huggingface.co/datasets/wikipedia) and my own json data) and use set_transform/with_transform for preprocessing. But it hangs at the end of the 1st epoch if dataloader_num_workers>=1. No problem with single worker.

## Steps to reproduce the bug
```python
train_dataset = datasets.load_dataset("wikipedia", "20220301.en",
split='train',
cache_dir=model_args.cache_dir,
streaming=False)
train_dataset.set_transform(psg_parse_fn)
train_dataloader = DataLoader(
train_dataset,
batch_size=args.train_batch_size,
sampler=DistributedSampler(train_dataset),
collate_fn=data_collator,
drop_last=args.dataloader_drop_last,
num_workers=args.dataloader_num_workers,
)
```

## Expected results

## Actual results
It simply hangs. The ending step is num_example/batch_size (one epoch).

## Environment info
- `datasets` version: 2.4.1.dev0
- Platform: Linux-5.4.170+-x86_64-with-glibc2.17
- Python version: 3.8.12
- PyArrow version: 8.0.0
- Pandas version: 1.4.1