Update handler.py

handler.py

@@ -66,11 +66,10 @@ class EndpointHandler():
                 model_output = self.onnx_extractor(sentence)  
             print("E")
 
-            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(max_pooled_output)
+            sentence_embeddings.append(max_pooling(model_output))
             print("G")
         return sentence_embeddings