correct model loading logic

app.py

@@ -26,7 +26,7 @@ from session import Session
 from model import load_model_and_dataloader, get_model_response
 
 # Load model and tokenizer once at the start
-model_path = "checkpoints/lora_grounded_obj_ref_checkpoint-4896"
+model_path = "checkpoints/llava_lora_grounded_obj_ref_checkpoint-4896"
 model_base = "checkpoints/llava-llama-2-7b-chat-lightning-preview"
 load_8bit = False
 load_4bit = False

llava/model/builder.py

@@ -197,7 +197,7 @@ def load_pretrained_model(
                 model_base,
                 torch_dtype=torch.bfloat16,
                 low_cpu_mem_usage=True,
-                device_map="auto",
+                device_map=device_map,
             )
             print(f"Loading LoRA weights from {model_path}")
             model = PeftModel.from_pretrained(model, model_path)