sarang-shrivastava commited on
Commit
1a9b1a2
1 Parent(s): 01e3d98

Update handler

Browse files
Files changed (1) hide show
  1. handler.py +2 -2
handler.py CHANGED
@@ -20,11 +20,11 @@ class EndpointHandler():
20
  def __init__(self, path=""):
21
 
22
  self.processor = Blip2Processor.from_pretrained(path)
23
- self.model = Blip2ForConditionalGeneration.from_pretrained(path, load_in_8bit=True, device_map="auto")
24
 
25
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
26
 
27
- # self.model.to(self.device)
28
 
29
  logging.info('Model moved to device-' + self.device)
30
 
 
20
  def __init__(self, path=""):
21
 
22
  self.processor = Blip2Processor.from_pretrained(path)
23
+ self.model = Blip2ForConditionalGeneration.from_pretrained(path, torch_dtype=torch.float16 , device_map="auto")
24
 
25
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
26
 
27
+ self.model.to(self.device)
28
 
29
  logging.info('Model moved to device-' + self.device)
30