Update modeling.py

modeling.py

@@ -59,6 +59,8 @@ class EmbeddingMixer(nn.Module):
         self.image_start_index = len(original_embedding)
 
     def set_image_embeddings(self, image_embeddings):
+        if len(image_embeddings.shape) == 3:
+            image_embeddings = image_embeddings.squeeze(0) # remove batch dim
         end_index = self.image_start_index + image_embeddings.shape[0]
         self.embedding[self.image_start_index:end_index] = image_embeddings