Update llava/model/builder.py

llava/model/builder.py

@@ -24,8 +24,7 @@ from llava.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, D
 
 
 def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cuda"):
-    # kwargs = {"device_map": device_map}
-    kwargs = {"device_map": device_map, "offload_folder": "offload"}
+    kwargs = {"device_map": device_map, "offload_folder": "offload", "load_in_8bit_fp32_cpu_offload": True}
 
     if device != "cuda":
         kwargs['device_map'] = {"": device}