Update libra/model/builder.py

libra/model/builder.py

@@ -27,7 +27,7 @@ def load_pretrained_model(model_path, model_base, model_name, device="cpu"):
     device_map = {"": device}
     kwargs = {
         "device_map": device_map,
-        "torch_dtype": torch.float32  # 对于 CPU，建议使用 float32 或 bfloat16
+        "torch_dtype": torch.bfloat16 
     }
 
         
@@ -114,7 +114,7 @@ def load_pretrained_model(model_path, model_base, model_name, device="cpu"):
         vision_tower = model.get_vision_tower()
         if not vision_tower.is_loaded:
             vision_tower.load_model()
-        vision_tower.to(device=device, dtype=torch.float32)
+        vision_tower.to(device=device, dtype=torch.bfloat16)
         image_processor = vision_tower.image_processor
 
     if hasattr(model.config, "max_sequence_length"):