Image dataset with_transform not applied

Beginners
1

Hi,

I am training a computer vision model and want to apply AutoImageProcessor to prepare images for model.

When I use with_transform and the use a trainer, transformation is not applied.

Dataset contains PIL images.

from transformers import AutoImageProcessor
image_processor = AutoImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k', use_fast=True)

def transform(example):
    ds = {}
    ds['image'] = image_processor(example, return_tensors='pt')['pixel_values'].reshape(3,224,224)
    return ds

dataset = dataset.with_transform(transform)

When I use .map it works, but I have a large dataset and map is taking too much space.

Any ideas why with_transform is not called at all? I also have tried with DataLoader and transformation is not applied.