Issue: Multiprocessing with `Dataset.map` and  `PyTorch` results in deadlock 

Comment:
Hi, I have the same problem in use **datasets.IterableDatasetDict.map()**
my pytorch is 2.0.0a0+gitc263bd4
my python is 3.8.16(default, Jun 12 2023, 17:37:21)
work on aarch64 in 16 node, each node with 4*nVidia-A100-40G
every node have 4 process execute code as ↓

```
from datasets import load_dataset, interleave_datasets, IterableDatasetDict, concatenate_datasets
```
...
```
    model_args.cache_dir = '/home/scx/.cache'
    for dataset_name in data_args.datasets_name:
        train_datasets.append(
            load_dataset(
                dataset_name,
                cache_dir=model_args.cache_dir,
                use_auth_token=True if model_args.use_auth_token else None,
                streaming=data_args.streaming,
                split='train'
            ).select_columns('text')
        )
        valid_datasets.append(
            load_dataset(
                dataset_name,
                cache_dir=model_args.cache_dir,
                use_auth_token=True if model_args.use_auth_token else None,
                streaming=data_args.streaming,
                split='validation'
            ).select_columns('text')
        )
        train_dataset = interleave_datasets(train_datasets,
                                        probabilities=data_args.datasets_probabilities, 
                                        seed=training_args.seed,
                                        stopping_strategy='all_exhausted')
    raw_datasets = IterableDatasetDict({'train': train_dataset, 'validation': valid_dataset})
```
...

```
    tokenized_datasets = None
    with training_args.main_process_first(desc="dataset map tokenization"):
        if not data_args.streaming:
            tokenized_datasets = raw_datasets.map(
                tokenize_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on dataset",
                remove_columns=column_names,
            )
        else:
            #TODO 20230722
            logger.info('{}: {}'.format(__file__, 'tokenized_datasets = raw_datasets.map('))
            logger.info('len raw_datasets: {}'.format(len(raw_datasets.items())))
            logger.info('raw_datasets:{}'.format(raw_datasets.items()))
            tokenized_datasets = raw_datasets.map(
                tokenize_function,
                batched=True,
                batch_size=1000,
                remove_columns=column_names
            )
            logger.info('map ok!')
            logger.info('show train: {}'.format(next(iter(tokenized_datasets['train']))))
            logger.info('ok')
            # ### RAW CODE ###
            # tokenized_datasets = raw_datasets.map(
            #     tokenize_function,
            #     batched=True,
            #     batch_size=1000,
            #     remove_columns=column_names
            # )
            #TODO 20230722
    logger.info("Finish tokenization")
```
the output of my code is
```
07/22/2023 21:57:09 - INFO - __main__ - /demo/run_blue_space.py: tokenized_datasets = raw_datasets.map(
07/22/2023 21:57:09 - INFO - __main__ - len raw_datasets: 2
07/22/2023 21:57:09 - INFO - __main__ - raw_datasets:dict_items([('train', <datasets.iterable_dataset.IterableDataset object at 0x4005ee301190>), ('validation', <datasets.iterable_dataset.IterableDataset object at 0x4005ee5427f0>)])
07/22/2023 21:57:09 - INFO - __main__ - map ok!
07/22/2023 22:01:07 - INFO - __main__ - show train: {'input_ids': [14608, 26797, 31891, 34260, 12227, 33207, 5, 5, 31632, 26797, 31891, 34260, 12227, 33207, 7398, 28561, 31236, 31177, 31253, 33558, 31556, 31377, 72, 20732, 32383, 32295, 14027, 31178, 53, 61, 53, 55, 31189, 31146, 31321, 31235, 53, 61, 56, 58, 31189, 31145, 72, 53, 61, 58, 54, 31189, 54, 31245, 53, 60, 31224, 31896, 31178, 28561, 29331, 20732, 31888, 32637, 4426, 2824, 72, 53, 61, 60, 55, 31189, 53, 54, 31245, 53, 31224, 31896, 31178, 28561, 29331, 26137, 20732, 4426, 2824, 73, 54, 52, 52, 52, 31189, 61, 31245, 59, 31224, 31896, 31178, 29331, 28561, 20732, 4426, 2824, 73, 5], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}
07/22/2023 22:01:07 - INFO - __main__ - ok
```