set_format does not work with multi dimension tensor
{'login': 'phihung', 'id': 5902432, 'node_id': 'MDQ6VXNlcjU5MDI0MzI=', 'avatar_url': 'https://avatars.githubusercontent.com/u/5902432?v=4', 'gravatar_id': '', 'url': 'https://api.github.com/users/phihung', 'html_url': 'https://github.com/phihung', 'followers_url': 'https://api.github.com/users/phihung/followers', 'following_url': 'https://api.github.com/users/phihung/following{/other_user}', 'gists_url': 'https://api.github.com/users/phihung/gists{/gist_id}', 'starred_url': 'https://api.github.com/users/phihung/starred{/owner}{/repo}', 'subscriptions_url': 'https://api.github.com/users/phihung/subscriptions', 'organizations_url': 'https://api.github.com/users/phihung/orgs', 'repos_url': 'https://api.github.com/users/phihung/repos', 'events_url': 'https://api.github.com/users/phihung/events{/privacy}', 'received_events_url': 'https://api.github.com/users/phihung/received_events', 'type': 'User', 'site_admin': False}
[{'id': 1935892857, 'node_id': 'MDU6TGFiZWwxOTM1ODkyODU3', 'url': 'https://api.github.com/repos/huggingface/datasets/labels/bug', 'name': 'bug', 'color': 'd73a4a', 'default': True, 'description': "Something isn't working"}]
[]
['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 :\r\n```python\r\nds = Dataset.from_dict({"A": [torch.rand((2, 2))]}, features=Features({"A": Array2D(shape=(2, 2), dtype="float32")}))\r\n```\r\n'
 '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`.\r\n```\r\ndataset = load_dataset("my_dataset.py")\r\ndataset = dataset.map(lambda example: blabla(example))\r\ndict_dataset_test = dataset["test"].to_dict()\r\n...\r\ndataset_test = Dataset.from_dict(dict_dataset_test, features=Features(features))\r\n```\r\nHowever, 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?'
 'Hi! Yes, you can directly pass the `Features` dictionary as `features` in `map` to cast the column to `Array2D`:\r\n```python\r\ndataset = dataset.map(lambda example: blabla(example), features=Features(features))\r\n```\r\nOr you can use `cast` after `map` to do that:\r\n```python\r\ndataset = dataset.map(lambda example: blabla(example))\r\ndataset = dataset.cast(Features(features))\r\n```'
 'Fantastic thank you @mariosasko\r\nThe first option you suggested is indeed way faster 😃 ']
## 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
{'url': 'https://api.github.com/repos/huggingface/datasets/issues/4007/reactions', 'total_count': 0, '+1': 0, '-1': 0, 'laugh': 0, 'hooray': 0, 'confused': 0, 'heart': 0, 'rocket': 0, 'eyes': 0}