minhdang commited on
Commit
82baabb
1 Parent(s): cf9f226

Update inference.py

Browse files
Files changed (1) hide show
  1. inference.py +2 -2
inference.py CHANGED
@@ -40,9 +40,9 @@ def load_model(model_path):
40
  vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
41
  tokenizer = vl_chat_processor.tokenizer
42
  vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
43
- model_path, trust_remote_code=True, load_in_8bit=True,low_cpu_mem_usage=True
44
  )
45
- vl_gpt = vl_gpt.cuda().eval()
46
  return tokenizer, vl_gpt, vl_chat_processor
47
 
48
 
 
40
  vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
41
  tokenizer = vl_chat_processor.tokenizer
42
  vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
43
+ model_path, trust_remote_code=True, load_in_8bit=True,low_cpu_mem_usage=True,device_map="auto"
44
  )
45
+ vl_gpt = vl_gpt.eval()
46
  return tokenizer, vl_gpt, vl_chat_processor
47
 
48