Shanshan Wang commited on
Commit
8ba3710
1 Parent(s): d259dc9
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -30,7 +30,8 @@ def load_model_and_set_image_function(model_name):
30
  torch_dtype=torch.bfloat16,
31
  low_cpu_mem_usage=True,
32
  trust_remote_code=True,
33
- use_auth_token=hf_token
 
34
  ).eval().cuda()
35
 
36
  tokenizer = AutoTokenizer.from_pretrained(
 
30
  torch_dtype=torch.bfloat16,
31
  low_cpu_mem_usage=True,
32
  trust_remote_code=True,
33
+ use_auth_token=hf_token,
34
+ device_map="auto"
35
  ).eval().cuda()
36
 
37
  tokenizer = AutoTokenizer.from_pretrained(