fix: txt_feats should repeat batch_size times

AILab-CVC · May 18, 2024 · 416fc7e · 416fc7e
1 parent 5c1e7b3
commit 416fc7e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/yolo_world/models/detectors/yolo_world.py b/yolo_world/models/detectors/yolo_world.py
@@ -79,6 +79,9 @@ def extract_feat(
         if batch_data_samples is None:
             texts = self.texts
             txt_feats = self.text_feats
+            batch_size=batch_inputs.shape[0]
+            texts = texts * batch_size
+            txt_feats = txt_feats.repeat(batch_size, 1, 1)
         elif isinstance(batch_data_samples,
                         dict) and 'texts' in batch_data_samples:
             texts = batch_data_samples['texts']