Update handler.py

handler.py

@@ -9,7 +9,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 class EndpointHandler():
     def __init__(self, path=""):
-        self.model = Swinv2Model.from_pretrained("microsoft/swinv2-large-patch4-window12to24-192to384-22kto1k-ft").to(device)
+        self.model = Swinv2Model.from_pretrained("microsoft/swinv2-large-patch4-window12to24-192to384-22kto1k-ft", add_pooling_layer = True).to(device)
         self.processor = AutoImageProcessor.from_pretrained("microsoft/swinv2-large-patch4-window12to24-192to384-22kto1k-ft")
 
     def __call__(self, data: Any) -> List[float]:
@@ -21,6 +21,6 @@ class EndpointHandler():
         with torch.no_grad():
             outputs = self.model(**inputs)
         
-        last_hidden_states = outputs.last_hidden_state
-        return last_hidden_states.tolist()
+        pooler_output = outputs.pooler_output
+        return pooler_output.tolist()