Update handler.py

handler.py

@@ -18,7 +18,9 @@ class EndpointHandler():
         # self.base_model = Blip2ForConditionalGeneration.from_pretrained(self.model_base, load_in_8bit=True)
         # self.pipe = Blip2ForConditionalGeneration.from_pretrained(self.model_base, load_in_8bit=True, torch_dtype=torch.float16)
         self.processor = Blip2Processor.from_pretrained(self.model_name)
-        self.model = BlipForQuestionAnswering.from_pretrained(self.model_name).to(self.device)
+        self.model = BlipForQuestionAnswering.from_pretrained(self.model_name,
+                                                              device_map="auto", load_in_8bit=True
+                                                             ).to(self.device)
         # self.model = PeftModel.from_pretrained(self.model_name, self.base_model_name).to(self.device)
 
         self.device = "cuda" if torch.cuda.is_available() else "cpu"