IterableDataset: Unsupported ScalarType BFloat16

[stoical07]

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

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:

{'x': tensor([1.], dtype=torch.bfloat16)}

Environment info

datasets==2.20.0
torch==2.2.2

[stoical07]

@lhoestq Thank you for merging #6607, but unfortunately the issue persists for IterableDataset 😔

[lhoestq]

Hi ! I opened #7002 to fix this bug

[stoical07]

Amazing, thank you so much @lhoestq! 🙏