Alright, I found it! Nothing to do with sizes at all! Turns out, formats of tensors were incorrect. I changed to:
self.lab[‘labels’] = torch.tensor(labels, dtype=torch.int64)
self.lab[‘boxes’] = torch.tensor(boxes, dtype = torch.float)
self.lab[‘masks’] = torch.tensor(masks, dtype = torch.uint8)
Before that they all were uint8. Damn it!