Update handler.py

handler.py

@@ -5,12 +5,9 @@ import base64
 from io import BytesIO
 from transformers import CLIPProcessor, CLIPModel
 
-# check for GPU
-device = 0 if torch.cuda.is_available() else -1
-
 class EndpointHandler():
     def __init__(self, path=""):
-        self.model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14-336")
+        self.model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14-336", device_map="auto")
         self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14-336")
 
     def __call__(self, data: Any) -> List[float]: