Update handler

handler.py

@@ -3,9 +3,10 @@ from typing import Dict, List, Any
 # from transformers import AutoTokenizer
 # import torch
 from datetime import datetime
+import torch
 
-
-
+import logging
+logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.DEBUG)
 
 import requests
 from PIL import Image
@@ -19,6 +20,12 @@ class EndpointHandler():
         self.processor = Blip2Processor.from_pretrained(path)
         self.model = Blip2ForConditionalGeneration.from_pretrained(path, device_map="auto")
 
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self.model.to(self.device)
+
+        logging.info('Model moved to device-' + self.device)
+        
         # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         # self.model.eval()
         # self.model.to(device=device, dtype=self.torch_dtype)
@@ -72,7 +79,7 @@ class EndpointHandler():
         raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
         
         question = "how many dogs are in the picture?"
-        inputs = self.processor(raw_image, question, return_tensors="pt").to("cuda")
+        inputs = self.processor(raw_image, question, return_tensors="pt").to(self.device)
         
         out = self.model.generate(**inputs)
         output_text = self.processor.decode(out[0], skip_special_tokens=True)