Update handler.py

handler.py

@@ -52,11 +52,11 @@ class EndpointHandler():
                 model_output = self.onnx_extractor(sentence)  
             print("E")
 
-            sentence_embeddings_max = [max(col) for col in zip(*model_output)]
+            max_pooled_output, _ = torch.max(model_output, dim=1)
             # Perform pooling. In this case, max pooling.    
             # embedding = mean_pooling(model_output, encoded_input['attention_mask'])
             print("F")
             
-            sentence_embeddings.append(sentence_embeddings_max)
+            sentence_embeddings.append(max_pooled_output)
             print("G")
         return sentence_embeddings