Value Error when trying to make Labels to tensor

Question

I am trying to get my custom dataloader working. I plan on using this as a yolo network. Right now, I am using transforms.v2 in pytorch for my transforms. With this, I use SanitizeBoundingBoxes() to clean up any invalid boxes. This should also remove the corresponding label to the bounding box (i think). I am passing the labels as a tuple with a nested dict.

When v2.SanitizeBoundingBoxes(labels_getter=None) the code works, but all the labels are loaded, even the ones that are not in the cropped image.

Removing the labels_getter=None, the code encounters an error: ValueError: The labels in the input to forward() must be a tensor or None, got <class 'list'> instead.

From here, I think that I should convert the labels (list: [0, 1, 0, 0, 2, 0, 1]) to a tensor. I use this line in the dataloader: img_labels = torch.from_numpy(np.asarray(img_labels)). This produces a ValueError, but I am not sure what kind, as it is blank. Ill show some of the output.

ValueError: Caught ValueError in DataLoader worker process 0.
.
.
.

File ~/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py:1520, in Module._call_impl(self, *args, **kwargs)
   1515 # If we don't have any hooks, we want to skip the rest of the logic in
   1516 # this function, and just call forward.
   1517 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1518         or _global_backward_pre_hooks or _global_backward_hooks
   1519         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1520     return forward_call(*args, **kwargs)
...
---> 46     raise ValueError
     47 elif tensor.ndim == 2:
     48     tensor = tensor.unsqueeze(0)

ValueError: 

I still feel it has something to do with the forward(), but I am not sure.

This is my custom Dataloader:

class AGR_Dataset(Dataset):
    def __init__(self, annotations_root, img_root, dataset_labels, transform=None):
        """
        Arguments:
            annotations_root (string): Path to the csv file with annotations.
            img_root (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.annotations_root = annotations_root
        self.img_root = img_root
        self.transform = transform
        self.dataset_labels = dataset_labels

    def __len__(self):
        return len(self.annotations_root)
    
    def __getitem__(self, idx):
        # idx may be the index or image name, I think image name
        if torch.is_tensor(idx):
            idx = idx.tolist()
        
        idx_name = os.listdir(self.img_root)[idx]
        
        img_name = os.path.join(self.img_root, idx_name)
        annotation_data = os.path.join(self.annotations_root, f"{idx_name.removesuffix('.jpg')}.txt")

        image = io.imread(img_name)
        with open(annotation_data, 'r') as file:
            lines = file.readlines()
            img_data = []
            img_labels = []
            for line in lines:
                line = line.split(',')
                line = [i.strip() for i in line]
                line = [float(num) for num in line[0].split()]
                img_labels.append(int(line[0]))
                img_data.append(line[1:])
                
        img_data = np.array(img_data)
        img_data = np.multiply(img_data, [image.shape[1], image.shape[0], image.shape[1], image.shape[0]])
        boxes = tv_tensors.BoundingBoxes(img_data, format='CXCYWH', canvas_size=(image.shape[0], image.shape[1]))

        img_labels = torch.from_numpy(np.asarray(img_labels))
        target = {'boxes': boxes, 'labels': img_labels}
        
        if self.transform:
            image, target = self.transform(image, target)
    
        return image, target

I load in my image path and annotation path (txt doc with bouding boxes, 1 txt file per image)

annotations_path = '7053442/Annotations/Yolo'
img_path = '7053442'
dataset_labels = ['car', 'bus', 'truck']

batchSize = 4
numWorkers = 4

I load my transforms and create the dataloader:

data_transform = v2.Compose([
    v2.ToImage(),
    v2.RandomResizedCrop(size=(680, 680), antialias=True),
    v2.ClampBoundingBoxes(),
    # v2.SanitizeBoundingBoxes(labels_getter=None),
    v2.SanitizeBoundingBoxes(),
    
])

testDataset = AGR_Dataset(f'{annotations_path}/test/', 
                        f'{img_path}/test/', dataset_labels,
                        transform=data_transform)

testDataloader = DataLoader(testDataset, collate_fn=collate_fn, batch_size=batchSize,
                        shuffle=False, num_workers=numWorkers)

These two commands give me the errors:

dl = iter(testDataloader)
for i in range(0, 4):
    sample_batched, target = next(dl)
    print(len(target[0]['boxes']))

img, target = testDataset[1]
print(f"{type(img) = }\n{type(target) = }\n{target.keys() = }")
print(f"{type(target['boxes']) = }\n{type(target['labels']) = }")
print(target['labels'])

I have tried taking everything out of the transform except the ToImage().