Print logits for debugging.

utils/model.py

@@ -458,6 +458,9 @@ class OwlViTForClassification(nn.Module):
                     loss_dict["loss_sym_box_label"] = sym_loss_box_label
                 # ----------------------------------------------------------------------------------------
 
+        #DEBUG:
+        print(f"im_features size: {image_feats.shape}, text_embeds size: {text_embeds.shape}")
+        print(f"im_features sum: {image_feats.sum().item()}, text_embeds sum: {text_embeds.sum().item()}")
         # Predict image-level classes (batch_size, num_patches, num_queries)
         image_text_logits, pred_logits, part_logits = self.cls_head(image_feats, text_embeds, topk_idxs)
 
@@ -475,6 +478,9 @@ class OwlViTForClassification(nn.Module):
                 sym_loss = self.loss_symmetric(logits_per_image, logits_per_text, targets_cls)
                 loss_dict["loss_xclip"] = sym_loss
 
+        #DEBUG:
+        print(f"pred_logits size: {part_logits.shape}, pred_logits size: {part_logits.shape}")
+        print(f"part_logits sum: {pred_logits.sum().item()}, part_logits sum: {pred_logits.sum().item()}")
         return pred_logits, part_logits, loss_dict
 
     def loss_symmetric(self, text_logits: torch.Tensor, image_logits: torch.Tensor, targets: torch.Tensor, box_labels: torch.Tensor = None) -> torch.Tensor: