Fixed normalization error when only 1 image is sent to endpoint

handler.py

@@ -219,13 +219,14 @@ class EndpointHandler:
                     # get image embeddings
                     batch_emb = self.model_clip.get_image_features(pixel_values=batch)
                     # detach text emb from graph, move to CPU, and convert to numpy array
-                    batch_emb = batch_emb.squeeze(0)
+                    # Check the shape of the tensor before squeezing
+                    if batch_emb.dim() > 1 and batch_emb.shape[0] == 1:
+                        batch_emb = batch_emb.squeeze(0)
+
                     batch_emb = batch_emb.cpu().detach().numpy()
                     # NORMALIZE
-                    if batch_emb.ndim == 1:
-                        batch_emb = batch_emb / np.linalg.norm(batch_emb, axis=0)
-                    else:
-                        batch_emb = batch_emb.T / np.linalg.norm(batch_emb, axis=1)
+
+                    batch_emb = batch_emb.T / np.linalg.norm(batch_emb, axis=1)
                     # transpose back to (21, 512)
                     batch_emb = batch_emb.T.tolist()
                 embedding_end_time = timeit.default_timer()