RajatMalviya commited on
Commit
6a1856b
·
verified ·
1 Parent(s): add552c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -49,7 +49,8 @@ def load_llava_model():
49
  llava_model = AutoModelForCausalLM.from_pretrained(
50
  model_id,
51
  torch_dtype=torch.float16,
52
- device_map="auto"
 
53
  )
54
 
55
  return llava_model, llava_processor
 
49
  llava_model = AutoModelForCausalLM.from_pretrained(
50
  model_id,
51
  torch_dtype=torch.float16,
52
+ device_map="auto",
53
+ cache_dir="./models" # Add this line
54
  )
55
 
56
  return llava_model, llava_processor