diff --git a/datadreamer/dataset_annotation/owlv2_annotator.py b/datadreamer/dataset_annotation/owlv2_annotator.py
index 6a2068a..8d9eeed 100644
--- a/datadreamer/dataset_annotation/owlv2_annotator.py
+++ b/datadreamer/dataset_annotation/owlv2_annotator.py
@@ -98,7 +98,7 @@ def _generate_annotations(
         """
         n = len(images)
         batched_prompts = [prompts] * n
-        target_sizes = torch.Tensor(images[0].size[::-1]).repeat((n, 1)).to(self.device)
+        target_sizes = torch.Tensor([img.size[::-1] for img in images]).to(self.device)
 
         # resize the images to the model's input size
         img_size = (1008, 1008) if self.size == "large" else (960, 960)