Setting to torch format not working with torchvision and MNIST
CONTRIBUTOR
Adding to the previous information, I think `torch.utils.data.DataLoader` is doing some conversion. 
What I tried:
```python
train_dataset = load_dataset('mnist')
```
I don't use any `map` or `set_format` or any `transform`. I use this directly, and try to load batches using the `DataLoader` with batch size 2, I get an output like this for the `image`:

```
[[tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor...
```
For `label`, it works fine:
```
tensor([7, 6])
```
Note that I didn't specify conversion to torch tensors anywhere.

Basically, there are two problems here:
1. `dataset.map` doesn't return tensor type objects, even though it uses the transforms, the grayscale conversion in transform was done, but the output was lists only.
2.  The `DataLoader` performs its own conversion, which may be not desired.

I understand that we can't change `DataLoader` because it is a torch functionality, however, is there a way we can handle image data to allow using it with torch `DataLoader` and `torchvision` properly?

I think if the `image` was a torch tensor (N,H,W,C), or  a list of torch tensors (H,W,C), before it is passed to `DataLoader`, then we might not face this issue. 
2021-03-08T09:11:33Z
https://github.com/huggingface/datasets/issues/2005#issuecomment-792603507
https://api.github.com/repos/huggingface/datasets/issues/2005
MDEyOklzc3VlQ29tbWVudDc5MjYwMzUwNw==
https://api.github.com/repos/huggingface/datasets/issues/comments/792603507/reactions
2021-03-08T10:49:25Z
https://api.github.com/repos/huggingface/datasets/issues/comments/792603507
https://avatars.githubusercontent.com/u/29076344?v=4
https://api.github.com/users/gchhablani/events{/privacy}
https://api.github.com/users/gchhablani/followers
https://api.github.com/users/gchhablani/following{/other_user}
https://api.github.com/users/gchhablani/gists{/gist_id}
https://github.com/gchhablani
gchhablani
MDQ6VXNlcjI5MDc2MzQ0
https://api.github.com/users/gchhablani/orgs
https://api.github.com/users/gchhablani/received_events
https://api.github.com/users/gchhablani/repos
https://api.github.com/users/gchhablani/starred{/owner}{/repo}
https://api.github.com/users/gchhablani/subscriptions
User
https://api.github.com/users/gchhablani
public
MEMBER
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
2021-03-09T10:59:22Z
https://github.com/huggingface/datasets/issues/2005#issuecomment-793716753
https://api.github.com/repos/huggingface/datasets/issues/2005
MDEyOklzc3VlQ29tbWVudDc5MzcxNjc1Mw==
https://api.github.com/repos/huggingface/datasets/issues/comments/793716753/reactions
2021-03-09T10:59:22Z
https://api.github.com/repos/huggingface/datasets/issues/comments/793716753
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
Hi @lhoestq

Raw feature types are like this:
```
Image:
<class 'list'> 60000  #(type, len)
<class 'list'> 28
<class 'list'> 28
<class 'int'>
Label:
<class 'list'> 60000
<class 'int'>
```
Inside the `prepare_feature` method with batch size 100000 , after processing, they are like this:

Inside Prepare Train Features
```
Image:
<class 'list'> 10000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'list'> 10000
<class 'torch.Tensor'>
```

After map, the feature type are like this:
```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'float'>
Label:
<class 'list'> 60000
<class 'int'>
```

After dataloader with batch size 2, the batch features are like this:
```
Image:
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```
<hr>

When I was setting the format of `train_dataset` to 'torch' after mapping  - 
```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'>
```

Corresponding DataLoader batch:
```
From DataLoader batch features
Image:
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```

I will check with features and get back.




2021-03-09T13:28:47Z
https://github.com/huggingface/datasets/issues/2005#issuecomment-793898873
https://api.github.com/repos/huggingface/datasets/issues/2005
MDEyOklzc3VlQ29tbWVudDc5Mzg5ODg3Mw==
https://api.github.com/repos/huggingface/datasets/issues/comments/793898873/reactions
2021-03-09T14:41:40Z
https://api.github.com/repos/huggingface/datasets/issues/comments/793898873
https://avatars.githubusercontent.com/u/29076344?v=4
https://api.github.com/users/gchhablani/events{/privacy}
https://api.github.com/users/gchhablani/followers
https://api.github.com/users/gchhablani/following{/other_user}
https://api.github.com/users/gchhablani/gists{/gist_id}
https://github.com/gchhablani
gchhablani
MDQ6VXNlcjI5MDc2MzQ0
https://api.github.com/users/gchhablani/orgs
https://api.github.com/users/gchhablani/received_events
https://api.github.com/users/gchhablani/repos
https://api.github.com/users/gchhablani/starred{/owner}{/repo}
https://api.github.com/users/gchhablani/subscriptions
User
https://api.github.com/users/gchhablani
public
CONTRIBUTOR
Hi @lhoestq

# Using Array3D
I tried this:
```python
features = datasets.Features({
            "image": datasets.Array3D(shape=(1,28,28),dtype="float32"),
            "label": datasets.features.ClassLabel(names=["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]),
        })
train_dataset = raw_dataset.map(prepare_features, features = features,batched=True, batch_size=10000)
```
and it didn't fix the issue.

During the `prepare_train_features:
```
Image:
<class 'list'> 10000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'list'> 10000
<class 'torch.Tensor'>
```

After the `map`:

```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'float'>
Label:
<class 'list'> 60000
<class 'int'>
```
From the DataLoader batch:
```
Image:
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```
It is the same as before.

---

Using `datasets.Sequence(datasets.Array2D(shape=(28,28),dtype="float32"))` gave an error during `map`:

```python
ArrowNotImplementedError                  Traceback (most recent call last)
<ipython-input-95-d28e69289084> in <module>()
      3             "label": datasets.features.ClassLabel(names=["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]),
      4         })
----> 5 train_dataset = raw_dataset.map(prepare_features, features = features,batched=True, batch_size=10000)

15 frames
/usr/local/lib/python3.7/dist-packages/datasets/dataset_dict.py in map(self, function, with_indices, input_columns, batched, batch_size, remove_columns, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc)
    446                     num_proc=num_proc,
    447                 )
--> 448                 for k, dataset in self.items()
    449             }
    450         )

/usr/local/lib/python3.7/dist-packages/datasets/dataset_dict.py in <dictcomp>(.0)
    446                     num_proc=num_proc,
    447                 )
--> 448                 for k, dataset in self.items()
    449             }
    450         )

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in map(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint)
   1307                 fn_kwargs=fn_kwargs,
   1308                 new_fingerprint=new_fingerprint,
-> 1309                 update_data=update_data,
   1310             )
   1311         else:

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    202         }
    203         # apply actual function
--> 204         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    205         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    206         # re-apply format to the output

/usr/local/lib/python3.7/dist-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    335             # Call actual function
    336 
--> 337             out = func(self, *args, **kwargs)
    338 
    339             # Update fingerprint of in-place transforms + update in-place history of transforms

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in _map_single(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset, update_data)
   1580                     if update_data:
   1581                         batch = cast_to_python_objects(batch)
-> 1582                         writer.write_batch(batch)
   1583             if update_data:
   1584                 writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file

/usr/local/lib/python3.7/dist-packages/datasets/arrow_writer.py in write_batch(self, batch_examples, writer_batch_size)
    274             typed_sequence = TypedSequence(batch_examples[col], type=col_type, try_type=col_try_type)
    275             typed_sequence_examples[col] = typed_sequence
--> 276         pa_table = pa.Table.from_pydict(typed_sequence_examples)
    277         self.write_table(pa_table, writer_batch_size)
    278 

/usr/local/lib/python3.7/dist-packages/pyarrow/table.pxi in pyarrow.lib.Table.from_pydict()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.asarray()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.array()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib._handle_arrow_array_protocol()

/usr/local/lib/python3.7/dist-packages/datasets/arrow_writer.py in __arrow_array__(self, type)
     95                 out = pa.ExtensionArray.from_storage(type, pa.array(self.data, type.storage_dtype))
     96             else:
---> 97                 out = pa.array(self.data, type=type)
     98             if trying_type and out[0].as_py() != self.data[0]:
     99                 raise TypeError(

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.array()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib._sequence_to_array()

/usr/local/lib/python3.7/dist-packages/pyarrow/error.pxi in pyarrow.lib.pyarrow_internal_check_status()

/usr/local/lib/python3.7/dist-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowNotImplementedError: extension
```
2021-03-09T14:38:15Z
https://github.com/huggingface/datasets/issues/2005#issuecomment-793981891
https://api.github.com/repos/huggingface/datasets/issues/2005
MDEyOklzc3VlQ29tbWVudDc5Mzk4MTg5MQ==
https://api.github.com/repos/huggingface/datasets/issues/comments/793981891/reactions
2021-03-09T15:21:28Z
https://api.github.com/repos/huggingface/datasets/issues/comments/793981891
https://avatars.githubusercontent.com/u/29076344?v=4
https://api.github.com/users/gchhablani/events{/privacy}
https://api.github.com/users/gchhablani/followers
https://api.github.com/users/gchhablani/following{/other_user}
https://api.github.com/users/gchhablani/gists{/gist_id}
https://github.com/gchhablani
gchhablani
MDQ6VXNlcjI5MDc2MzQ0
https://api.github.com/users/gchhablani/orgs
https://api.github.com/users/gchhablani/received_events
https://api.github.com/users/gchhablani/repos
https://api.github.com/users/gchhablani/starred{/owner}{/repo}
https://api.github.com/users/gchhablani/subscriptions
User
https://api.github.com/users/gchhablani
public
CONTRIBUTOR
# Convert raw tensors to torch format
Strangely, converting to torch tensors works perfectly on `raw_dataset`:
```python
raw_dataset.set_format('torch',columns=['image','label'])
```
Types:
```
Image:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'>
```

Using this for transforms:
```python
def prepare_features(examples):
    images = []
    labels = []
    for example_idx, example in enumerate(examples["image"]):
        if transform is not None:
            images.append(transform(
                examples["image"][example_idx].numpy()
            ))
        else:
            images.append(examples["image"][example_idx].numpy())
        labels.append(examples["label"][example_idx])
    output = {"label":labels, "image":images}
    return output
```

Inside `prepare_train_features`:
```
Image:
<class 'list'> 10000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'list'> 10000
<class 'torch.Tensor'>
```

After `map`:
```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'>
```
DataLoader batch:

```
Image:
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```

---

## Using `torch` format:
```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'>
```
DataLoader batches:

```
Image:
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```

---
## Using the features - `Array3D`:

```
Image:
<class 'list'> 10000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'list'> 10000
<class 'torch.Tensor'>
```

After `map`:
```
Image:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'>
```

After DataLoader `batch`:
```
Image:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```

The last one works perfectly.

![image](https://user-images.githubusercontent.com/29076344/110491452-4cf09c00-8117-11eb-8a47-73bf3fc0c3dc.png)

I wonder why this worked, and others didn't.







2021-03-09T14:54:38Z
https://github.com/huggingface/datasets/issues/2005#issuecomment-794003895
https://api.github.com/repos/huggingface/datasets/issues/2005
MDEyOklzc3VlQ29tbWVudDc5NDAwMzg5NQ==
https://api.github.com/repos/huggingface/datasets/issues/comments/794003895/reactions
2021-03-09T15:20:46Z
https://api.github.com/repos/huggingface/datasets/issues/comments/794003895
https://avatars.githubusercontent.com/u/29076344?v=4
https://api.github.com/users/gchhablani/events{/privacy}
https://api.github.com/users/gchhablani/followers
https://api.github.com/users/gchhablani/following{/other_user}
https://api.github.com/users/gchhablani/gists{/gist_id}
https://github.com/gchhablani
gchhablani
MDQ6VXNlcjI5MDc2MzQ0
https://api.github.com/users/gchhablani/orgs
https://api.github.com/users/gchhablani/received_events
https://api.github.com/users/gchhablani/repos
https://api.github.com/users/gchhablani/starred{/owner}{/repo}
https://api.github.com/users/gchhablani/subscriptions
User
https://api.github.com/users/gchhablani
public
CONTRIBUTOR
Concluding, the way it works right now is:

1. Converting raw dataset to `torch` format.
2. Use the transform and apply using `map`, ensure the returned values are tensors. 
3. When mapping, use `features` with `image` being `Array3D` type.
2021-03-09T15:13:56Z
https://github.com/huggingface/datasets/issues/2005#issuecomment-794028636
https://api.github.com/repos/huggingface/datasets/issues/2005
MDEyOklzc3VlQ29tbWVudDc5NDAyODYzNg==
https://api.github.com/repos/huggingface/datasets/issues/comments/794028636/reactions
2021-03-09T15:14:44Z
https://api.github.com/repos/huggingface/datasets/issues/comments/794028636
https://avatars.githubusercontent.com/u/29076344?v=4
https://api.github.com/users/gchhablani/events{/privacy}
https://api.github.com/users/gchhablani/followers
https://api.github.com/users/gchhablani/following{/other_user}
https://api.github.com/users/gchhablani/gists{/gist_id}
https://github.com/gchhablani
gchhablani
MDQ6VXNlcjI5MDc2MzQ0
https://api.github.com/users/gchhablani/orgs
https://api.github.com/users/gchhablani/received_events
https://api.github.com/users/gchhablani/repos
https://api.github.com/users/gchhablani/starred{/owner}{/repo}
https://api.github.com/users/gchhablani/subscriptions
User
https://api.github.com/users/gchhablani
public
MEMBER
What the dataset returns depends on the feature type.
For a feature type that is Sequence(Sequence(Sequence(Value("uint8")))), a dataset formatted as "torch" return lists of lists of tensors. This is because the lists lengths may vary.
For a feature type that is Array3D on the other hand it returns one tensor. This is because the size of the tensor is fixed and defined bu the Array3D type.
2021-03-09T15:41:58Z
https://github.com/huggingface/datasets/issues/2005#issuecomment-794064284
https://api.github.com/repos/huggingface/datasets/issues/2005
MDEyOklzc3VlQ29tbWVudDc5NDA2NDI4NA==
https://api.github.com/repos/huggingface/datasets/issues/comments/794064284/reactions
2021-03-09T15:41:58Z
https://api.github.com/repos/huggingface/datasets/issues/comments/794064284
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
Okay, that makes sense.
Raw images are list of Array2D, hence we get a single tensor when `set_format` is used. But, why should I need to convert the raw images to `torch` format when `map` does this internally?

Using `Array3D` did not work with `map` when raw images weren't `set_format`ted to torch type.
2021-03-09T15:49:32Z
https://github.com/huggingface/datasets/issues/2005#issuecomment-794073691
https://api.github.com/repos/huggingface/datasets/issues/2005
MDEyOklzc3VlQ29tbWVudDc5NDA3MzY5MQ==
https://api.github.com/repos/huggingface/datasets/issues/comments/794073691/reactions
2021-03-09T15:49:32Z
https://api.github.com/repos/huggingface/datasets/issues/comments/794073691
https://avatars.githubusercontent.com/u/29076344?v=4
https://api.github.com/users/gchhablani/events{/privacy}
https://api.github.com/users/gchhablani/followers
https://api.github.com/users/gchhablani/following{/other_user}
https://api.github.com/users/gchhablani/gists{/gist_id}
https://github.com/gchhablani
gchhablani
MDQ6VXNlcjI5MDc2MzQ0
https://api.github.com/users/gchhablani/orgs
https://api.github.com/users/gchhablani/received_events
https://api.github.com/users/gchhablani/repos
https://api.github.com/users/gchhablani/starred{/owner}{/repo}
https://api.github.com/users/gchhablani/subscriptions
User
https://api.github.com/users/gchhablani
public
CONTRIBUTOR
I understand that `map` needs to know what kind of output tensors are expected, and thus converting the raw dataset to `torch` format is necessary. Closing the issue since it is resolved.
2021-03-09T17:58:06Z
https://github.com/huggingface/datasets/issues/2005#issuecomment-794236789
https://api.github.com/repos/huggingface/datasets/issues/2005
MDEyOklzc3VlQ29tbWVudDc5NDIzNjc4OQ==
https://api.github.com/repos/huggingface/datasets/issues/comments/794236789/reactions
2021-03-09T17:58:06Z
https://api.github.com/repos/huggingface/datasets/issues/comments/794236789
https://avatars.githubusercontent.com/u/29076344?v=4
https://api.github.com/users/gchhablani/events{/privacy}
https://api.github.com/users/gchhablani/followers
https://api.github.com/users/gchhablani/following{/other_user}
https://api.github.com/users/gchhablani/gists{/gist_id}
https://github.com/gchhablani
gchhablani
MDQ6VXNlcjI5MDc2MzQ0
https://api.github.com/users/gchhablani/orgs
https://api.github.com/users/gchhablani/received_events
https://api.github.com/users/gchhablani/repos
https://api.github.com/users/gchhablani/starred{/owner}{/repo}
https://api.github.com/users/gchhablani/subscriptions
User
https://api.github.com/users/gchhablani
public
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.