Using `set_format(type='torch', columns=columns)` makes Array2D/3D columns stop formatting correctly
I have a Dataset with two Features defined as follows:

```
'image': Array3D(dtype="int64", shape=(3, 224, 224)),
'bbox': Array2D(dtype="int64", shape=(512, 4)),
```

On said dataset, if I `dataset.set_format(type='torch')` and then use the dataset in a dataloader, these columns are correctly cast to Tensors of (batch_size, 3, 224, 244) for example.

However, if I `dataset.set_format(type='torch', columns=['image', 'bbox'])` these columns are cast to Lists of tensors and miss the batch size completely (the 3 dimension is the list length).

I'm currently digging through datasets formatting code to try and find out why, but was curious if someone knew an immediate solution for this.