vit-gpt2-coco-en Lora Fine-Tuning error

[DanielPuentee]

I am trying to fine-tune a Vit model for an image-to-text problem.

The error is the following:

ValueError: You have to specify pixel_values
The error seems clear, however, I can't seem to fix it. The training and validation set already has the value of pixel_values

Code:

class Dataset(torch.utils.data.Dataset):
    def __init__(self, labels, pixel_values):
        self.labels = list(labels)
        self.pixel_values = list(pixel_values)

    def __len__(self):
        return len(self.pixel_values)

    def __getitem__(self, idx):
        pixel_values = self.pixel_values[idx]
        labels = self.labels[idx]
        return {"labels": labels, "pixel_values": pixel_values}

dataset_train = Dataset(input_pixel_values_train, output_prompt_ids_train) 
dataset_val = Dataset(input_pixel_values_val, output_prompt_ids_val)

I think the problem comes in the construction of the LoraConfig in particular in the task_type parameter.

config = LoraConfig(
    r=8, 
    lora_alpha=32,
    target_modules=["query","value"],
    lora_dropout=0.05,
    bias="none",
    task_type="SEQ_CLS" 
)

I have tried all the possibilities, which I think they are: SEQ_CLS, TOKEN_CLS, SEQ_2_SEQ_LM & CAUSAL_LM. I don't know if that's what the error is, but that's where I think it's going.

So if anyone knows how to fix it or if the error is somewhere else, please tell me.