apoorvkh
/

blip2-flan-t5-xxl-endpoint

apoorvkh commited on Jun 18, 2023

Commit

76aff17

•

1 Parent(s): 665939a

Adding torch_dtype = float16

Files changed (1) hide show

handler.py CHANGED Viewed

@@ -22,7 +22,7 @@ class EndpointHandler():
         self.model = Blip2ForConditionalGeneration.from_pretrained(
             "Salesforce/blip2-flan-t5-xxl", device_map=device_map,
-            # torch_dtype=torch.float16
             load_in_8bit=True, load_in_8bit_fp32_cpu_offload=True
         )

         self.model = Blip2ForConditionalGeneration.from_pretrained(
             "Salesforce/blip2-flan-t5-xxl", device_map=device_map,
+            torch_dtype=torch.float16,
             load_in_8bit=True, load_in_8bit_fp32_cpu_offload=True
         )