John6666 commited on
Commit
3733821
·
verified ·
1 Parent(s): 625b8ab

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -250,7 +250,7 @@ def load_models():
250
  assert isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)), f"Tokenizer is of type {type(tokenizer)}"
251
 
252
  print(f"Loading LLM: {MODEL_PATH} 🤖")
253
- text_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map=device, torch_dtype=torch.bfloat16).eval() # device_map=auto may cause LoRA error
254
 
255
  if LORA_PATH.exists():
256
  print("Loading VLM's custom text model 🤖")
 
250
  assert isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)), f"Tokenizer is of type {type(tokenizer)}"
251
 
252
  print(f"Loading LLM: {MODEL_PATH} 🤖")
253
+ text_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto", torch_dtype=torch.bfloat16).eval()
254
 
255
  if LORA_PATH.exists():
256
  print("Loading VLM's custom text model 🤖")