Training hangs at the end of epoch, with set_transform/with_transform+multiple workers
{'login': 'memray', 'id': 4197249, 'node_id': 'MDQ6VXNlcjQxOTcyNDk=', 'avatar_url': 'https://avatars.githubusercontent.com/u/4197249?v=4', 'gravatar_id': '', 'url': 'https://api.github.com/users/memray', 'html_url': 'https://github.com/memray', 'followers_url': 'https://api.github.com/users/memray/followers', 'following_url': 'https://api.github.com/users/memray/following{/other_user}', 'gists_url': 'https://api.github.com/users/memray/gists{/gist_id}', 'starred_url': 'https://api.github.com/users/memray/starred{/owner}{/repo}', 'subscriptions_url': 'https://api.github.com/users/memray/subscriptions', 'organizations_url': 'https://api.github.com/users/memray/orgs', 'repos_url': 'https://api.github.com/users/memray/repos', 'events_url': 'https://api.github.com/users/memray/events{/privacy}', 'received_events_url': 'https://api.github.com/users/memray/received_events', 'type': 'User', 'site_admin': False}
[{'id': 1935892857, 'node_id': 'MDU6TGFiZWwxOTM1ODkyODU3', 'url': 'https://api.github.com/repos/huggingface/datasets/labels/bug', 'name': 'bug', 'color': 'd73a4a', 'default': True, 'description': "Something isn't working"}]
[]
## 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
{'url': 'https://api.github.com/repos/huggingface/datasets/issues/4774/reactions', 'total_count': 0, '+1': 0, '-1': 0, 'laugh': 0, 'hooray': 0, 'confused': 0, 'heart': 0, 'rocket': 0, 'eyes': 0}