Update libra/model/builder.py

libra/model/builder.py

@@ -23,12 +23,18 @@ from libra.model import *
 from libra.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
 
 
-def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cuda"):
+def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cup"):
     kwargs = {"device_map": device_map}
 
+    
     if device != "cuda":
         kwargs['device_map'] = {"": device}
         
+    if device == "cpu":
+        kwargs["torch_dtype"] = torch.float32
+    else:
+        kwargs["torch_dtype"] = torch.float16
+        
     if load_8bit:
         kwargs['load_in_8bit'] = True
     elif load_4bit: