Setting to torch format not working with torchvision and MNIST
What's the feature types of your new dataset after `.map` ?

Can you try with adding `features=` in the `.map` call in order to set the "image" feature type to `Array2D` ?
The default feature type is lists of lists, we've not implemented shape verification to use ArrayXD instead of nested lists yet
Hi

I am trying to use `torchvision.transforms` to handle the transformation of the image data in the `mnist` dataset. Assume I have a `transform` variable which contains the `torchvision.transforms` object.

A snippet of what I am trying to do:
```python
def prepare_features(examples):
    images = []
    labels = []
    for example_idx, example in enumerate(examples["image"]):
        if transform is not None:
            images.append(transform(
                np.array(examples["image"][example_idx], dtype=np.uint8)
            ))
        else:
            images.append(torch.tensor(np.array(examples["image"][example_idx], dtype=np.uint8)))
        labels.append(torch.tensor(examples["label"][example_idx]))
    output = {"label":labels, "image":images}
    return output

raw_dataset = load_dataset('mnist')
train_dataset = raw_dataset.map(prepare_features, batched=True, batch_size=10000)
train_dataset.set_format("torch",columns=["image","label"])
```

After this, I check the type of the following:
```python
print(type(train_dataset["train"]["label"]))
print(type(train_dataset["train"]["image"][0]))
```
This leads to the following output:

```python
<class 'torch.Tensor'>
<class 'list'>
```
I use `torch.utils.DataLoader` for batches, the type of `batch["train"]["image"]` is also `<class 'list'>`.

I don't understand why only the `label` is converted to a torch tensor, why does the image not get converted? How can I fix this issue?

Thanks,
Gunjan

EDIT:
I just checked the shapes, and the types, `batch[image]` is a actually a list of list of tensors. Shape is (1,28,2,28), where `batch_size` is 2. I don't understand why this is happening. Ideally it should be a tensor of shape (2,1,28,28).

EDIT 2:
Inside `prepare_train_features`, the shape of `images[0]` is `torch.Size([1,28,28])`, the conversion is working. However, the output of the `map` is a list of list of list of list.
Setting to torch format not working with torchvision and MNIST 
 Hi

I am trying to use `torchvision.transforms` to handle the transformation of the image data in the `mnist` dataset. Assume I have a `transform` variable which contains the `torchvision.transforms` object.

A snippet of what I am trying to do:
```python
def prepare_features(examples):
    images = []
    labels = []
    for example_idx, example in enumerate(examples["image"]):
        if transform is not None:
            images.append(transform(
                np.array(examples["image"][example_idx], dtype=np.uint8)
            ))
        else:
            images.append(torch.tensor(np.array(examples["image"][example_idx], dtype=np.uint8)))
        labels.append(torch.tensor(examples["label"][example_idx]))
    output = {"label":labels, "image":images}
    return output

raw_dataset = load_dataset('mnist')
train_dataset = raw_dataset.map(prepare_features, batched=True, batch_size=10000)
train_dataset.set_format("torch",columns=["image","label"])
```

After this, I check the type of the following:
```python
print(type(train_dataset["train"]["label"]))
print(type(train_dataset["train"]["image"][0]))
```
This leads to the following output:

```python
<class 'torch.Tensor'>
<class 'list'>
```
I use `torch.utils.DataLoader` for batches, the type of `batch["train"]["image"]` is also `<class 'list'>`.

I don't understand why only the `label` is converted to a torch tensor, why does the image not get converted? How can I fix this issue?

Thanks,
Gunjan

EDIT:
I just checked the shapes, and the types, `batch[image]` is a actually a list of list of tensors. Shape is (1,28,2,28), where `batch_size` is 2. I don't understand why this is happening. Ideally it should be a tensor of shape (2,1,28,28).

EDIT 2:
Inside `prepare_train_features`, the shape of `images[0]` is `torch.Size([1,28,28])`, the conversion is working. However, the output of the `map` is a list of list of list of list. 
 What's the feature types of your new dataset after `.map` ?

Can you try with adding `features=` in the `.map` call in order to set the "image" feature type to `Array2D` ?
The default feature type is lists of lists, we've not implemented shape verification to use ArrayXD instead of nested lists yet