set_format does not work with multi dimension tensor
COLLABORATOR
Hi! Use the `ArrayXD` feature type (where X is the number of dimensions) to get correctly formated tensors. So in your case, define the dataset as follows :
```python
ds = Dataset.from_dict({"A": [torch.rand((2, 2))]}, features=Features({"A": Array2D(shape=(2, 2), dtype="float32")}))
```

2022-03-24T13:41:44Z
https://github.com/huggingface/datasets/issues/4007#issuecomment-1077641791
https://api.github.com/repos/huggingface/datasets/issues/4007
IC_kwDODunzps5AO4I_
https://api.github.com/repos/huggingface/datasets/issues/comments/1077641791/reactions
2022-03-24T13:41:44Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1077641791
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
NONE
Hi @mariosasko I'm facing the same issue and the only work around I've found so far is to convert my `DatasetDict` to a dictionary and then create new objects with `Dataset.from_dict`.
```
dataset = load_dataset("my_dataset.py")
dataset = dataset.map(lambda example: blabla(example))
dict_dataset_test = dataset["test"].to_dict()
...
dataset_test = Dataset.from_dict(dict_dataset_test, features=Features(features))
```
However, converting a `Dataset` object to a dict takes quite a lot of time and memory... Is there a way to directly create an `Array2D` without having to transform the original `Dataset` to a dict?
2022-03-29T14:40:08Z
https://github.com/huggingface/datasets/issues/4007#issuecomment-1081956678
https://api.github.com/repos/huggingface/datasets/issues/4007
IC_kwDODunzps5AfVlG
https://api.github.com/repos/huggingface/datasets/issues/comments/1081956678/reactions
2022-03-29T14:40:08Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1081956678
https://avatars.githubusercontent.com/u/32683010?v=4
https://api.github.com/users/JulesBelveze/events{/privacy}
https://api.github.com/users/JulesBelveze/followers
https://api.github.com/users/JulesBelveze/following{/other_user}
https://api.github.com/users/JulesBelveze/gists{/gist_id}
https://github.com/JulesBelveze
JulesBelveze
MDQ6VXNlcjMyNjgzMDEw
https://api.github.com/users/JulesBelveze/orgs
https://api.github.com/users/JulesBelveze/received_events
https://api.github.com/users/JulesBelveze/repos
https://api.github.com/users/JulesBelveze/starred{/owner}{/repo}
https://api.github.com/users/JulesBelveze/subscriptions
User
https://api.github.com/users/JulesBelveze
public
COLLABORATOR
Hi! Yes, you can directly pass the `Features` dictionary as `features` in `map` to cast the column to `Array2D`:
```python
dataset = dataset.map(lambda example: blabla(example), features=Features(features))
```
Or you can use `cast` after `map` to do that:
```python
dataset = dataset.map(lambda example: blabla(example))
dataset = dataset.cast(Features(features))
```
2022-03-29T15:28:17Z
https://github.com/huggingface/datasets/issues/4007#issuecomment-1082018126
https://api.github.com/repos/huggingface/datasets/issues/4007
IC_kwDODunzps5AfklO
https://api.github.com/repos/huggingface/datasets/issues/comments/1082018126/reactions
2022-03-29T15:28:17Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1082018126
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
NONE
Fantastic thank you @mariosasko
The first option you suggested is indeed way faster 😃 
2022-03-30T07:12:01Z
https://github.com/huggingface/datasets/issues/4007#issuecomment-1082710887
https://api.github.com/repos/huggingface/datasets/issues/4007
IC_kwDODunzps5AiNtn
https://api.github.com/repos/huggingface/datasets/issues/comments/1082710887/reactions
2022-03-30T07:28:57Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1082710887
https://avatars.githubusercontent.com/u/32683010?v=4
https://api.github.com/users/JulesBelveze/events{/privacy}
https://api.github.com/users/JulesBelveze/followers
https://api.github.com/users/JulesBelveze/following{/other_user}
https://api.github.com/users/JulesBelveze/gists{/gist_id}
https://github.com/JulesBelveze
JulesBelveze
MDQ6VXNlcjMyNjgzMDEw
https://api.github.com/users/JulesBelveze/orgs
https://api.github.com/users/JulesBelveze/received_events
https://api.github.com/users/JulesBelveze/repos
https://api.github.com/users/JulesBelveze/starred{/owner}{/repo}
https://api.github.com/users/JulesBelveze/subscriptions
User
https://api.github.com/users/JulesBelveze
public
## Describe the bug
set_format only transforms the last dimension of a multi-dimension list to tensor

## Steps to reproduce the bug
```python
import torch
from datasets import Dataset

ds = Dataset.from_dict({"A": [torch.rand((2, 2))]})
# ds = Dataset.from_dict({"A": [np.random.rand(2, 2)]})  # => same result
ds = ds.with_format("torch")
print(ds[0])
```

## Expected results
```
{'A': [tensor([[0.6689, 0.1516], [0.1403, 0.5567]])]}
```

## Actual results
```
{'A': [tensor([0.6689, 0.1516]), tensor([0.1403, 0.5567])]}
```

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- datasets version: 2.0.0
- Platform: Mac OSX
- Python version: 3.8.12
- PyArrow version: 7.0.0