Update handler.py

handler.py

@@ -11,7 +11,7 @@ import base64
         
 class EndpointHandler():
     def __init__(self, path=""):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
         self.model_base = "Salesforce/blip2-opt-2.7b"
         self.model_name = "sooh-j/blip2-vizwizqa"
         # self.pipe = Blip2ForConditionalGeneration.from_pretrained(self.model_base, load_in_8bit=True, torch_dtype=torch.float16)
@@ -101,7 +101,7 @@ class EndpointHandler():
         # )
     
         with torch.no_grad():
-            out = self.model.generate(**processed)
+            out = self.model.generate(**processed).to(self.device)
         
         result = {}
         text_output = self.processor.decode(out[0], skip_special_tokens=True)