IterableDataset: Unsupported ScalarType BFloat16
NONE
@lhoestq Thank you for merging #6607, but unfortunately the issue persists for `IterableDataset` :pensive: 
2024-06-25T14:54:44Z
https://github.com/huggingface/datasets/issues/7000#issuecomment-2189182668
https://api.github.com/repos/huggingface/datasets/issues/7000
IC_kwDODunzps6CfEbM
https://api.github.com/repos/huggingface/datasets/issues/comments/2189182668/reactions
2024-06-25T14:54:44Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2189182668
https://avatars.githubusercontent.com/u/170015089?v=4
https://api.github.com/users/stoical07/events{/privacy}
https://api.github.com/users/stoical07/followers
https://api.github.com/users/stoical07/following{/other_user}
https://api.github.com/users/stoical07/gists{/gist_id}
https://github.com/stoical07
stoical07
U_kgDOCiI5cQ
https://api.github.com/users/stoical07/orgs
https://api.github.com/users/stoical07/received_events
https://api.github.com/users/stoical07/repos
https://api.github.com/users/stoical07/starred{/owner}{/repo}
https://api.github.com/users/stoical07/subscriptions
User
https://api.github.com/users/stoical07
public
MEMBER
Hi ! I opened https://github.com/huggingface/datasets/pull/7002 to fix this bug
2024-06-25T15:38:24Z
https://github.com/huggingface/datasets/issues/7000#issuecomment-2189290471
https://api.github.com/repos/huggingface/datasets/issues/7000
IC_kwDODunzps6Cfevn
https://api.github.com/repos/huggingface/datasets/issues/comments/2189290471/reactions
2024-06-25T15:38:24Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2189290471
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
Amazing, thank you so much @lhoestq! :pray:
2024-06-25T16:03:59Z
https://github.com/huggingface/datasets/issues/7000#issuecomment-2189358244
https://api.github.com/repos/huggingface/datasets/issues/7000
IC_kwDODunzps6CfvSk
https://api.github.com/repos/huggingface/datasets/issues/comments/2189358244/reactions
2024-06-25T16:03:59Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2189358244
https://avatars.githubusercontent.com/u/170015089?v=4
https://api.github.com/users/stoical07/events{/privacy}
https://api.github.com/users/stoical07/followers
https://api.github.com/users/stoical07/following{/other_user}
https://api.github.com/users/stoical07/gists{/gist_id}
https://github.com/stoical07
stoical07
U_kgDOCiI5cQ
https://api.github.com/users/stoical07/orgs
https://api.github.com/users/stoical07/received_events
https://api.github.com/users/stoical07/repos
https://api.github.com/users/stoical07/starred{/owner}{/repo}
https://api.github.com/users/stoical07/subscriptions
User
https://api.github.com/users/stoical07
public
### Describe the bug

`IterableDataset.from_generator` crashes when using BFloat16:

```
  File "/usr/local/lib/python3.11/site-packages/datasets/utils/_dill.py", line 169, in _save_torchTensor
    args = (obj.detach().cpu().numpy(),)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^

TypeError: Got unsupported ScalarType BFloat16
```

### Steps to reproduce the bug

```python
import torch
from datasets import IterableDataset


def demo(x):
    yield {"x": x}


x = torch.tensor([1.], dtype=torch.bfloat16)

dataset = IterableDataset.from_generator(
    demo,
    gen_kwargs=dict(x=x),
)

example = next(iter(dataset))
print(example)
```

### Expected behavior

Code sample should print:
```python
{'x': tensor([1.], dtype=torch.bfloat16)}
```

### Environment info

```
datasets==2.20.0
torch==2.2.2
```