feat

utils.py

@@ -55,7 +55,10 @@ def post_process_grounded_object_detection_with_queries(
     Now also returns the query embeddings for each kept prediction.
     """
     logits, boxes = outputs.logits, outputs.pred_boxes
-    assert len(logits) == queries.shape[0], "logits and queries must have the same batch size, but got {} and {}".format(len(logits), queries.shape[0])
+    print("logits: ", logits.shape)
+    print("boxes: ", boxes.shape)
+    print("queries: ", queries.shape)
+    assert len(logits[0]) == queries.shape[0], "logits and queries must have the same batch size, but got {} and {}".format(len(logits), queries.shape[0])
 
     probs = torch.sigmoid(logits)  # (batch_size, num_queries, 256)
     scores = torch.max(probs, dim=-1)[0]  # (batch_size, num_queries)