Setting to torch format not working with torchvision and MNIST
Hi @lhoestq

# Using Array3D
I tried this:
```python
features = datasets.Features({
            "image": datasets.Array3D(shape=(1,28,28),dtype="float32"),
            "label": datasets.features.ClassLabel(names=["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]),
        })
train_dataset = raw_dataset.map(prepare_features, features = features,batched=True, batch_size=10000)
```
and it didn't fix the issue.

During the `prepare_train_features:
```
Image:
<class 'list'> 10000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'list'> 10000
<class 'torch.Tensor'>
```

After the `map`:

```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'float'>
Label:
<class 'list'> 60000
<class 'int'>
```
From the DataLoader batch:
```
Image:
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```
It is the same as before.

---

Using `datasets.Sequence(datasets.Array2D(shape=(28,28),dtype="float32"))` gave an error during `map`:

```python
ArrowNotImplementedError                  Traceback (most recent call last)
<ipython-input-95-d28e69289084> in <module>()
      3             "label": datasets.features.ClassLabel(names=["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]),
      4         })
----> 5 train_dataset = raw_dataset.map(prepare_features, features = features,batched=True, batch_size=10000)

15 frames
/usr/local/lib/python3.7/dist-packages/datasets/dataset_dict.py in map(self, function, with_indices, input_columns, batched, batch_size, remove_columns, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc)
    446                     num_proc=num_proc,
    447                 )
--> 448                 for k, dataset in self.items()
    449             }
    450         )

/usr/local/lib/python3.7/dist-packages/datasets/dataset_dict.py in <dictcomp>(.0)
    446                     num_proc=num_proc,
    447                 )
--> 448                 for k, dataset in self.items()
    449             }
    450         )

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in map(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint)
   1307                 fn_kwargs=fn_kwargs,
   1308                 new_fingerprint=new_fingerprint,
-> 1309                 update_data=update_data,
   1310             )
   1311         else:

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    202         }
    203         # apply actual function
--> 204         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    205         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    206         # re-apply format to the output

/usr/local/lib/python3.7/dist-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    335             # Call actual function
    336 
--> 337             out = func(self, *args, **kwargs)
    338 
    339             # Update fingerprint of in-place transforms + update in-place history of transforms

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in _map_single(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset, update_data)
   1580                     if update_data:
   1581                         batch = cast_to_python_objects(batch)
-> 1582                         writer.write_batch(batch)
   1583             if update_data:
   1584                 writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file

/usr/local/lib/python3.7/dist-packages/datasets/arrow_writer.py in write_batch(self, batch_examples, writer_batch_size)
    274             typed_sequence = TypedSequence(batch_examples[col], type=col_type, try_type=col_try_type)
    275             typed_sequence_examples[col] = typed_sequence
--> 276         pa_table = pa.Table.from_pydict(typed_sequence_examples)
    277         self.write_table(pa_table, writer_batch_size)
    278 

/usr/local/lib/python3.7/dist-packages/pyarrow/table.pxi in pyarrow.lib.Table.from_pydict()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.asarray()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.array()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib._handle_arrow_array_protocol()

/usr/local/lib/python3.7/dist-packages/datasets/arrow_writer.py in __arrow_array__(self, type)
     95                 out = pa.ExtensionArray.from_storage(type, pa.array(self.data, type.storage_dtype))
     96             else:
---> 97                 out = pa.array(self.data, type=type)
     98             if trying_type and out[0].as_py() != self.data[0]:
     99                 raise TypeError(

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.array()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib._sequence_to_array()

/usr/local/lib/python3.7/dist-packages/pyarrow/error.pxi in pyarrow.lib.pyarrow_internal_check_status()

/usr/local/lib/python3.7/dist-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowNotImplementedError: extension
```
Hi

I am trying to use `torchvision.transforms` to handle the transformation of the image data in the `mnist` dataset. Assume I have a `transform` variable which contains the `torchvision.transforms` object.

A snippet of what I am trying to do:
```python
def prepare_features(examples):
    images = []
    labels = []
    for example_idx, example in enumerate(examples["image"]):
        if transform is not None:
            images.append(transform(
                np.array(examples["image"][example_idx], dtype=np.uint8)
            ))
        else:
            images.append(torch.tensor(np.array(examples["image"][example_idx], dtype=np.uint8)))
        labels.append(torch.tensor(examples["label"][example_idx]))
    output = {"label":labels, "image":images}
    return output

raw_dataset = load_dataset('mnist')
train_dataset = raw_dataset.map(prepare_features, batched=True, batch_size=10000)
train_dataset.set_format("torch",columns=["image","label"])
```

After this, I check the type of the following:
```python
print(type(train_dataset["train"]["label"]))
print(type(train_dataset["train"]["image"][0]))
```
This leads to the following output:

```python
<class 'torch.Tensor'>
<class 'list'>
```
I use `torch.utils.DataLoader` for batches, the type of `batch["train"]["image"]` is also `<class 'list'>`.

I don't understand why only the `label` is converted to a torch tensor, why does the image not get converted? How can I fix this issue?

Thanks,
Gunjan

EDIT:
I just checked the shapes, and the types, `batch[image]` is a actually a list of list of tensors. Shape is (1,28,2,28), where `batch_size` is 2. I don't understand why this is happening. Ideally it should be a tensor of shape (2,1,28,28).

EDIT 2:
Inside `prepare_train_features`, the shape of `images[0]` is `torch.Size([1,28,28])`, the conversion is working. However, the output of the `map` is a list of list of list of list.
Setting to torch format not working with torchvision and MNIST 
 Hi

I am trying to use `torchvision.transforms` to handle the transformation of the image data in the `mnist` dataset. Assume I have a `transform` variable which contains the `torchvision.transforms` object.

A snippet of what I am trying to do:
```python
def prepare_features(examples):
    images = []
    labels = []
    for example_idx, example in enumerate(examples["image"]):
        if transform is not None:
            images.append(transform(
                np.array(examples["image"][example_idx], dtype=np.uint8)
            ))
        else:
            images.append(torch.tensor(np.array(examples["image"][example_idx], dtype=np.uint8)))
        labels.append(torch.tensor(examples["label"][example_idx]))
    output = {"label":labels, "image":images}
    return output

raw_dataset = load_dataset('mnist')
train_dataset = raw_dataset.map(prepare_features, batched=True, batch_size=10000)
train_dataset.set_format("torch",columns=["image","label"])
```

After this, I check the type of the following:
```python
print(type(train_dataset["train"]["label"]))
print(type(train_dataset["train"]["image"][0]))
```
This leads to the following output:

```python
<class 'torch.Tensor'>
<class 'list'>
```
I use `torch.utils.DataLoader` for batches, the type of `batch["train"]["image"]` is also `<class 'list'>`.

I don't understand why only the `label` is converted to a torch tensor, why does the image not get converted? How can I fix this issue?

Thanks,
Gunjan

EDIT:
I just checked the shapes, and the types, `batch[image]` is a actually a list of list of tensors. Shape is (1,28,2,28), where `batch_size` is 2. I don't understand why this is happening. Ideally it should be a tensor of shape (2,1,28,28).

EDIT 2:
Inside `prepare_train_features`, the shape of `images[0]` is `torch.Size([1,28,28])`, the conversion is working. However, the output of the `map` is a list of list of list of list. 
 Hi @lhoestq

# Using Array3D
I tried this:
```python
features = datasets.Features({
            "image": datasets.Array3D(shape=(1,28,28),dtype="float32"),
            "label": datasets.features.ClassLabel(names=["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]),
        })
train_dataset = raw_dataset.map(prepare_features, features = features,batched=True, batch_size=10000)
```
and it didn't fix the issue.

During the `prepare_train_features:
```
Image:
<class 'list'> 10000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'list'> 10000
<class 'torch.Tensor'>
```

After the `map`:

```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'float'>
Label:
<class 'list'> 60000
<class 'int'>
```
From the DataLoader batch:
```
Image:
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```
It is the same as before.

---

Using `datasets.Sequence(datasets.Array2D(shape=(28,28),dtype="float32"))` gave an error during `map`:

```python
ArrowNotImplementedError                  Traceback (most recent call last)
<ipython-input-95-d28e69289084> in <module>()
      3             "label": datasets.features.ClassLabel(names=["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]),
      4         })
----> 5 train_dataset = raw_dataset.map(prepare_features, features = features,batched=True, batch_size=10000)

15 frames
/usr/local/lib/python3.7/dist-packages/datasets/dataset_dict.py in map(self, function, with_indices, input_columns, batched, batch_size, remove_columns, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc)
    446                     num_proc=num_proc,
    447                 )
--> 448                 for k, dataset in self.items()
    449             }
    450         )

/usr/local/lib/python3.7/dist-packages/datasets/dataset_dict.py in <dictcomp>(.0)
    446                     num_proc=num_proc,
    447                 )
--> 448                 for k, dataset in self.items()
    449             }
    450         )

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in map(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint)
   1307                 fn_kwargs=fn_kwargs,
   1308                 new_fingerprint=new_fingerprint,
-> 1309                 update_data=update_data,
   1310             )
   1311         else:

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    202         }
    203         # apply actual function
--> 204         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    205         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    206         # re-apply format to the output

/usr/local/lib/python3.7/dist-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    335             # Call actual function
    336 
--> 337             out = func(self, *args, **kwargs)
    338 
    339             # Update fingerprint of in-place transforms + update in-place history of transforms

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in _map_single(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset, update_data)
   1580                     if update_data:
   1581                         batch = cast_to_python_objects(batch)
-> 1582                         writer.write_batch(batch)
   1583             if update_data:
   1584                 writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file

/usr/local/lib/python3.7/dist-packages/datasets/arrow_writer.py in write_batch(self, batch_examples, writer_batch_size)
    274             typed_sequence = TypedSequence(batch_examples[col], type=col_type, try_type=col_try_type)
    275             typed_sequence_examples[col] = typed_sequence
--> 276         pa_table = pa.Table.from_pydict(typed_sequence_examples)
    277         self.write_table(pa_table, writer_batch_size)
    278 

/usr/local/lib/python3.7/dist-packages/pyarrow/table.pxi in pyarrow.lib.Table.from_pydict()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.asarray()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.array()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib._handle_arrow_array_protocol()

/usr/local/lib/python3.7/dist-packages/datasets/arrow_writer.py in __arrow_array__(self, type)
     95                 out = pa.ExtensionArray.from_storage(type, pa.array(self.data, type.storage_dtype))
     96             else:
---> 97                 out = pa.array(self.data, type=type)
     98             if trying_type and out[0].as_py() != self.data[0]:
     99                 raise TypeError(

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.array()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib._sequence_to_array()

/usr/local/lib/python3.7/dist-packages/pyarrow/error.pxi in pyarrow.lib.pyarrow_internal_check_status()

/usr/local/lib/python3.7/dist-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowNotImplementedError: extension
```