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