Getting a dimension mismatch when training dataset on openAI CLIP

Question

Here's the code I've written:

  from PIL import Image
device = "cuda:0" if torch.cuda.is_available() else "cpu" # If using GPU then use mixed precision training.

class image_title_dataset(Dataset):
    def __init__(self, list_image_path, list_txt, max_caption_len=77):
        self.image_path = list_image_path
        self.tokenizer = clip.tokenize
        self.max_caption_len = max_caption_len
        self.captions = [self.tokenize_and_pad(caption) for caption in list_txt]

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

    def __getitem__(self, idx):
          image = preprocess(Image.open(self.image_path[idx]))
          caption = self.captions[idx]
          caption = caption.to(device)
          image = image.to(device)
          return image, caption

    def tokenize_and_pad(self, caption):
        tokenized_caption = self.tokenizer(caption)
        caption_len = tokenized_caption.shape[0]
        if caption_len < self.max_caption_len:
            padding = torch.zeros((self.max_caption_len - caption_len, *tokenized_caption.shape[1:]), dtype=tokenized_caption.dtype, device=tokenized_caption.device)
            tokenized_caption = torch.cat([tokenized_caption, padding], dim=0)
        else:
            tokenized_caption = tokenized_caption[:self.max_caption_len]
        return tokenized_caption

This is where I defined the inputs to the model

And below is my training loop:

loss_img = nn.CrossEntropyLoss()
loss_txt = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=5e-5,betas=(0.9,0.98),eps=1e-6,weight_decay=0.2)

    for epoch in range(2):
  for batch in train_dataloader:
      optimizer.zero_grad()

      images, texts = batch 
    
      images = images.to(device)
      texts = texts.to(device)
    
      logits_per_image, logits_per_text = model(images, texts)

      ground_truth = torch.arange(len(images), dtype=torch.long, device=device)
      ground_truth = ground_truth.unsqueeze(0)  # add a dimension for batch size
      
      total_loss = (loss_img(logits_per_image, ground_truth) + loss_txt(logits_per_text, ground_truth)) / 2
      total_loss.backward()
      
      if device == "cpu":
         optimizer.step()
      else:
         convert_models_to_fp32(model)
         optimizer.step()
         clip.model.convert_weights(model)

I'm getting this error: permute(sparse_coo): number of dimensions in the tensor input does not match the length of the desired ordering of dimensions i.e. input.dim() = 4 is not equal to len(dims) = 3

Here's how the text list looks like: [['pictures are in wall', 'white curtain', 'rug', 'wooden shelf', 'toys'], ['frame', 'cabinets'], ['sun flowers', 'walls', 'coffee maker on counter', 'stove']....] Each image has multiple phrases with it, and I'm working on a segmentation task.