Merge pull request haotian-liu#851 from l-salewski/kwargs-for-model-l…

…oading 🩹 make ``load_pretrained_model`` accept kwargs
dr-data · Nov 24, 2023 · 2ca20de · 2ca20de
2 parents 414cebd + 8fa7f58
commit 2ca20de
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llava/model/builder.py b/llava/model/builder.py
@@ -23,8 +23,8 @@
 from llava.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"):
-    kwargs = {"device_map": device_map}
+def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cuda", **kwargs):
+    kwargs = {"device_map": device_map, **kwargs}
 
     if device != "cuda":
         kwargs['device_map'] = {"": device}