Multiprocessing with `Dataset.map` and  `PyTorch` results in deadlock 
Labels: name: bug
description: Something isn't working

## Describe the bug
There seems to be an issue with using multiprocessing with `datasets.Dataset.map` (i.e. setting `num_proc` to a value greater than one) combined with a function that uses `torch` under the hood.  The subprocesses that `datasets.Dataset.map` spawns [a this step](https://github.com/huggingface/datasets/blob/1b935dab9d2f171a8c6294269421fe967eb55e34/src/datasets/arrow_dataset.py#L2663) go into wait mode forever.  

## Steps to reproduce the bug
The below code goes into  deadlock when `NUMBER_OF_PROCESSES` is greater than one.
```python

NUMBER_OF_PROCESSES = 2

from transformers import AutoTokenizer, AutoModel
from datasets import load_dataset

dataset = load_dataset("glue", "mrpc", split="train")
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model.to("cpu")

def cls_pooling(model_output):
    return model_output.last_hidden_state[:, 0]

def generate_embeddings_batched(examples):
    sentences_batch = list(examples['sentence1'])
    encoded_input = tokenizer(
        sentences_batch, padding=True, truncation=True, return_tensors="pt"
    )
    encoded_input = {k: v.to("cpu") for k, v in encoded_input.items()}
    model_output = model(**encoded_input)
    embeddings = cls_pooling(model_output)
    examples['embeddings'] = embeddings.detach().cpu().numpy() # 64, 384
    return examples

embeddings_dataset = dataset.map(
    generate_embeddings_batched, 
    batched=True,
    batch_size=10,
    num_proc=NUMBER_OF_PROCESSES
)
```
While debugging it I've seen that it gets "stuck" when calling `torch.nn.Embedding.forward` but some testing shows that the same happens with other functions from `torch.nn`. 

## Environment info
- Platform: Linux-5.14.0-1052-oem-x86_64-with-glibc2.31
- Python version: 3.9.14
- PyArrow version: 9.0.0
- Pandas version: 1.5.0

Not sure if this is a HF problem, a PyTorch problem or something I'm doing wrong..

Thanks!


Comments:
Hi ! Could it be an Out of Memory issue that could have killed one of the processes ? can you check your memory ?

Hi! I don't think it is a memory issue. I'm monitoring the main and spawn python processes and threads with `htop` and the memory does not peak. Besides, the example I've posted above should not be that demanding in terms of memory, right? (I have 32GB of RAM). 

Indeed it should be fine. I couldn't reproduce the error though - I ran your script on my side and it works fine. What version of pytorch are you using ?

Interesting.. I'm using `torch 1.12.1`

I also tried on colab and it works fine 🤔 
Maybe something is wrong with your installation of pytorch ?

Oh actually I just saw that you're using python 3.9

This could be related to https://github.com/huggingface/datasets/issues/4113

We'll fix that as soon as we can, in the meantime you can try to use use single process, or use an older version of python maybe ?

I tried with python 3.7 and the issue persists. In collab, which also uses 3.7 I don't get the issue, so yes I guess is something on mu side... will post it here if I manage to fix it

Hi! Which version of transformers are you using? I test the code on Colab (so python 3.7) with transformers 4.23.1, torch 1.12.1 and pyarrow 9.0.0 (also 6.x), it worked without stuck.