Issue: Setting to torch format not working with torchvision and MNIST

Comment:
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
```