Update utils/model.py

utils/model.py

@@ -421,9 +421,11 @@ class OwlViTForClassification(nn.Module):
             print(f"text_inputs_parts - input_ids: {text_inputs_parts['input_ids'].shape}. attention_mask : {text_inputs_parts['attention_mask'].shape}")
             seq_length = text_inputs_parts['input_ids'].shape[-1]
             position_ids = self.owlvit.text_model.embeddings.position_ids[:, :seq_length]
+            txt_embeds = self.owlvit.text_model.embeddings.token_embedding(text_inputs_parts['input_ids'])
             print(f"position_embedding: {self.owlvit.text_model.embeddings.position_embedding(position_ids).shape}")
-            print(f"text_embeds: {self.owlvit.text_model.embeddings.token_embedding(text_inputs_parts['input_ids']).shape}")
-            text_embeds_parts = self.owlvit.text_model.text_model.get_text_features(**text_inputs_parts)
+            print(f"text_embeds: {txt_embeds.shape}")
+            print(f"pos + emb: {(txt_embeds + position_ids).shape}")
+            text_embeds_parts = self.owlvit.text_model.get_text_features(**text_inputs_parts)
             
             # # Embed images and text queries
             query_mask, text_embeds_parts = self._get_text_query_mask(text_inputs_parts, text_embeds_parts, batch_size)