Issue: Mapping gets stuck at 99%

Comment:
I think `batch_size/writer_batch_size` is already at its lowest in the 2nd `.map` since `batched=False` implies `batch_size=1` and `len(ds) = 1000 = writer_batch_size`.

Here is also a bunch of stack traces when I interrupted the process:

<details>
  <summary>stack trace 1</summary>

```python
(pyg)[d623204@rosetta-bigviz01 stage-laurent-f]$ python src/random_scripts/uses_random_data.py 
Found cached dataset random_data (/local_scratch/lfainsin/.cache/huggingface/datasets/random_data/default/0.0.0/444e214e1d0e6298cfd3f2368323ec37073dc1439f618e19395b1f421c69b066)
Applying mean/std:  97%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████     | 967/1000 [00:01<00:00, 534.87 examples/s]Traceback (most recent call last):                                                                                                                                                                                      
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 179, in __arrow_array__
    storage = to_pyarrow_listarray(data, pa_type)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 1466, in to_pyarrow_listarray
    return pa.array(data, pa_type.storage_dtype)
  File "pyarrow/array.pxi", line 320, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 39, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 144, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 123, in pyarrow.lib.check_status
pyarrow.lib.ArrowTypeError: Could not convert tensor([[-1.0273, -0.8037, -0.6860],
        [-0.5034, -1.2685, -0.0558],
        [-1.0908, -1.1820, -0.3178],
        ...,
        [-0.8171,  0.1781, -0.5903],
        [ 0.4370,  1.9305,  0.5899],
        [-0.1426,  0.9053, -1.7559]]) with type Tensor: was not a sequence or recognized null for conversion to list type

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 3449, in _map_single
    writer.write(example)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 490, in write
    self.write_examples_on_file()
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 448, in write_examples_on_file
    self.write_batch(batch_examples=batch_examples)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 553, in write_batch
    arrays.append(pa.array(typed_sequence))
  File "pyarrow/array.pxi", line 236, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 110, in pyarrow.lib._handle_arrow_array_protocol
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 223, in __arrow_array__
    return pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 446, in cast_to_python_objects
    return _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 407, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 408, in <listcomp>
    _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 319, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 320, in <listcomp>
    _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 263, in _cast_to_python_objects
    def _cast_to_python_objects(obj: Any, only_1d_for_numpy: bool, optimize_list_casting: bool) -> Tuple[Any, bool]:
KeyboardInterrupt

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 179, in __arrow_array__
    storage = to_pyarrow_listarray(data, pa_type)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 1466, in to_pyarrow_listarray
    return pa.array(data, pa_type.storage_dtype)
  File "pyarrow/array.pxi", line 320, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 39, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 144, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 123, in pyarrow.lib.check_status
pyarrow.lib.ArrowTypeError: Could not convert tensor([[-1.0273, -0.8037, -0.6860],
        [-0.5034, -1.2685, -0.0558],
        [-1.0908, -1.1820, -0.3178],
        ...,
        [-0.8171,  0.1781, -0.5903],
        [ 0.4370,  1.9305,  0.5899],
        [-0.1426,  0.9053, -1.7559]]) with type Tensor: was not a sequence or recognized null for conversion to list type

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/gpfs_new/data/users/lfainsin/stage-laurent-f/src/random_scripts/uses_random_data.py", line 62, in <module>
    ds_normalized = ds.map(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 580, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 545, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 3087, in map
    for rank, done, content in Dataset._map_single(**dataset_kwargs):
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 3492, in _map_single
    writer.finalize()
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 584, in finalize
    self.write_examples_on_file()
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 448, in write_examples_on_file
    self.write_batch(batch_examples=batch_examples)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 553, in write_batch
    arrays.append(pa.array(typed_sequence))
  File "pyarrow/array.pxi", line 236, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 110, in pyarrow.lib._handle_arrow_array_protocol
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 223, in __arrow_array__
    return pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 446, in cast_to_python_objects
    return _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 407, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 408, in <listcomp>
    _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 319, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 319, in <listcomp>
    [
KeyboardInterrupt
```

</details>

<details>
  <summary>stack trace 2</summary>

```python
(pyg)[d623204@rosetta-bigviz01 stage-laurent-f]$ python src/random_scripts/uses_random_data.py 
Found cached dataset random_data (/local_scratch/lfainsin/.cache/huggingface/datasets/random_data/default/0.0.0/444e214e1d0e6298cfd3f2368323ec37073dc1439f618e19395b1f421c69b066)
Applying mean/std:  99%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏ | 988/1000 [00:20<00:00, 526.19 examples/s]Applying mean/std: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▊| 999/1000 [00:21<00:00,  9.66 examples/s]Traceback (most recent call last):                                                                                                                                                                                      
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 179, in __arrow_array__
    storage = to_pyarrow_listarray(data, pa_type)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 1466, in to_pyarrow_listarray
    return pa.array(data, pa_type.storage_dtype)
  File "pyarrow/array.pxi", line 320, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 39, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 144, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 123, in pyarrow.lib.check_status
pyarrow.lib.ArrowTypeError: Could not convert tensor([[-1.0273, -0.8037, -0.6860],
        [-0.5034, -1.2685, -0.0558],
        [-1.0908, -1.1820, -0.3178],
        ...,
        [-0.8171,  0.1781, -0.5903],
        [ 0.4370,  1.9305,  0.5899],
        [-0.1426,  0.9053, -1.7559]]) with type Tensor: was not a sequence or recognized null for conversion to list type

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 3449, in _map_single
    writer.write(example)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 490, in write
    self.write_examples_on_file()
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 448, in write_examples_on_file
    self.write_batch(batch_examples=batch_examples)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 553, in write_batch
    arrays.append(pa.array(typed_sequence))
  File "pyarrow/array.pxi", line 236, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 110, in pyarrow.lib._handle_arrow_array_protocol
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 223, in __arrow_array__
    return pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 446, in cast_to_python_objects
    return _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 407, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 408, in <listcomp>
    _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 319, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 320, in <listcomp>
    _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 263, in _cast_to_python_objects
    def _cast_to_python_objects(obj: Any, only_1d_for_numpy: bool, optimize_list_casting: bool) -> Tuple[Any, bool]:
KeyboardInterrupt

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 179, in __arrow_array__
    storage = to_pyarrow_listarray(data, pa_type)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 1466, in to_pyarrow_listarray
    return pa.array(data, pa_type.storage_dtype)
  File "pyarrow/array.pxi", line 320, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 39, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 144, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 123, in pyarrow.lib.check_status
pyarrow.lib.ArrowTypeError: Could not convert tensor([[-1.0273, -0.8037, -0.6860],
        [-0.5034, -1.2685, -0.0558],
        [-1.0908, -1.1820, -0.3178],
        ...,
        [-0.8171,  0.1781, -0.5903],
        [ 0.4370,  1.9305,  0.5899],
        [-0.1426,  0.9053, -1.7559]]) with type Tensor: was not a sequence or recognized null for conversion to list type

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/gpfs_new/data/users/lfainsin/stage-laurent-f/src/random_scripts/uses_random_data.py", line 62, in <module>
    ds_normalized = ds.map(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 580, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 545, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 3087, in map
    for rank, done, content in Dataset._map_single(**dataset_kwargs):
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 3492, in _map_single
    writer.finalize()
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 584, in finalize
    self.write_examples_on_file()
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 448, in write_examples_on_file
    self.write_batch(batch_examples=batch_examples)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 553, in write_batch
    arrays.append(pa.array(typed_sequence))
  File "pyarrow/array.pxi", line 236, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 110, in pyarrow.lib._handle_arrow_array_protocol
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 223, in __arrow_array__
    return pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 446, in cast_to_python_objects
    return _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 407, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 408, in <listcomp>
    _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 319, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 320, in <listcomp>
    _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 291, in _cast_to_python_objects
    if config.JAX_AVAILABLE and "jax" in sys.modules:
KeyboardInterrupt
```

</details>

<details>
  <summary>stack trace 3</summary>

```python
(pyg)[d623204@rosetta-bigviz01 stage-laurent-f]$ python src/random_scripts/uses_random_data.py 
Found cached dataset random_data (/local_scratch/lfainsin/.cache/huggingface/datasets/random_data/default/0.0.0/444e214e1d0e6298cfd3f2368323ec37073dc1439f618e19395b1f421c69b066)
Applying mean/std:  99%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▎ | 989/1000 [00:01<00:00, 504.80 examples/s]Traceback (most recent call last):                                                                                                                                                                                      
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 179, in __arrow_array__
    storage = to_pyarrow_listarray(data, pa_type)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 1466, in to_pyarrow_listarray
    return pa.array(data, pa_type.storage_dtype)
  File "pyarrow/array.pxi", line 320, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 39, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 144, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 123, in pyarrow.lib.check_status
pyarrow.lib.ArrowTypeError: Could not convert tensor([[-1.0273, -0.8037, -0.6860],
        [-0.5034, -1.2685, -0.0558],
        [-1.0908, -1.1820, -0.3178],
        ...,
        [-0.8171,  0.1781, -0.5903],
        [ 0.4370,  1.9305,  0.5899],
        [-0.1426,  0.9053, -1.7559]]) with type Tensor: was not a sequence or recognized null for conversion to list type

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 3449, in _map_single
    writer.write(example)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 490, in write
    self.write_examples_on_file()
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 448, in write_examples_on_file
    self.write_batch(batch_examples=batch_examples)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 553, in write_batch
    arrays.append(pa.array(typed_sequence))
  File "pyarrow/array.pxi", line 236, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 110, in pyarrow.lib._handle_arrow_array_protocol
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 223, in __arrow_array__
    return pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 446, in cast_to_python_objects
    return _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 407, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 408, in <listcomp>
    _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 319, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 320, in <listcomp>
    _cast_to_python_objects(
KeyboardInterrupt

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 179, in __arrow_array__
    storage = to_pyarrow_listarray(data, pa_type)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 1466, in to_pyarrow_listarray
    return pa.array(data, pa_type.storage_dtype)
  File "pyarrow/array.pxi", line 320, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 39, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 144, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 123, in pyarrow.lib.check_status
pyarrow.lib.ArrowTypeError: Could not convert tensor([[-1.0273, -0.8037, -0.6860],
        [-0.5034, -1.2685, -0.0558],
        [-1.0908, -1.1820, -0.3178],
        ...,
        [-0.8171,  0.1781, -0.5903],
        [ 0.4370,  1.9305,  0.5899],
        [-0.1426,  0.9053, -1.7559]]) with type Tensor: was not a sequence or recognized null for conversion to list type

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/gpfs_new/data/users/lfainsin/stage-laurent-f/src/random_scripts/uses_random_data.py", line 62, in <module>
    ds_normalized = ds.map(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 580, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 545, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 3087, in map
    for rank, done, content in Dataset._map_single(**dataset_kwargs):
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 3492, in _map_single
    writer.finalize()
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 584, in finalize
    self.write_examples_on_file()
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 448, in write_examples_on_file
    self.write_batch(batch_examples=batch_examples)
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 553, in write_batch
    arrays.append(pa.array(typed_sequence))
  File "pyarrow/array.pxi", line 236, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 110, in pyarrow.lib._handle_arrow_array_protocol
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/arrow_writer.py", line 223, in __arrow_array__
    return pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 446, in cast_to_python_objects
    return _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 407, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 408, in <listcomp>
    _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 319, in _cast_to_python_objects
    [
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 320, in <listcomp>
    _cast_to_python_objects(
  File "/local_scratch/lfainsin/.conda/envs/pyg/lib/python3.10/site-packages/datasets/features/features.py", line 298, in _cast_to_python_objects
    if obj.ndim == 0:
KeyboardInterrupt
```

</details>