qnguyen3 commited on
Commit
13b8f1b
·
verified ·
1 Parent(s): a2b8740

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -21,7 +21,8 @@ model = LlavaQwen2ForCausalLM.from_pretrained(
21
  'qnguyen3/nanoLLaVA-1.5',
22
  torch_dtype=torch.float16,
23
  attn_implementation="flash_attention_2",
24
- trust_remote_code=True)
 
25
 
26
  class KeywordsStoppingCriteria(StoppingCriteria):
27
  def __init__(self, keywords, tokenizer, input_ids):
 
21
  'qnguyen3/nanoLLaVA-1.5',
22
  torch_dtype=torch.float16,
23
  attn_implementation="flash_attention_2",
24
+ trust_remote_code=True,
25
+ device_map='cpu')
26
 
27
  class KeywordsStoppingCriteria(StoppingCriteria):
28
  def __init__(self, keywords, tokenizer, input_ids):