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
-                    # Check the shape of the tensor before squeezing
-                    # if batch_emb.dim() > 1 and batch_emb.shape[0] == 1:
+
                     self.logger.info(
                         f"Shape of batch_emb after get_image_features: {batch_emb.shape}"
                     )
 
-                    batch_emb = batch_emb.squeeze(0)
+                    # Check the shape of the tensor before squeezing
+                    if batch_emb.shape[0] > 1:
+                        batch_emb = batch_emb.squeeze(0)
                     self.logger.info(
                         f"Shape of batch_emb after squeeze: {batch_emb.shape}"
                     )
@@ -239,9 +240,7 @@ class EndpointHandler:
                         )
                         # transpose back to (21, 512)
                         batch_emb = batch_emb.T.tolist()
-                        self.logger.info(
-                            f"Shape of batch_emb after transpose (2D case): {batch_emb.shape}"
-                        )
+
                 embedding_end_time = timeit.default_timer()
                 self.logger.info(
                     f"Embedding calculation took {embedding_end_time - embedding_start_time} seconds"