Update modeling.py

modeling.py

@@ -59,8 +59,7 @@ 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
+        assert len(image_embeddings.shape) == 2, "image_embeddings should be 2D"
         end_index = self.image_start_index + image_embeddings.shape[0]
         self.embedding[self.image_start_index:end_index] = image_embeddings