set_format does not work with multi dimension tensor
## 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