zhengr commited on
Commit
c675aaf
·
verified ·
1 Parent(s): 19e2804

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -30,10 +30,11 @@ streamer = TextStreamer(tokenizer)
30
  #model = AutoModelForCausalLM.from_pretrained(model, trust_remote_code=True, token = access_token, device_map="cuda", load_in_8bit=True, offload_folder="offload/")
31
  #model = PeftModel.from_pretrained(model, peft_model, offload_folder="offload/")
32
 
33
- model = AutoModelForCausalLM.from_pretrained(model, trust_remote_code=True, device_map="cuda")
34
  model = PeftModel.from_pretrained(model, peft_model)
35
 
36
- device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
37
  model.to(device)
38
 
39
  model = model.eval()
 
30
  #model = AutoModelForCausalLM.from_pretrained(model, trust_remote_code=True, token = access_token, device_map="cuda", load_in_8bit=True, offload_folder="offload/")
31
  #model = PeftModel.from_pretrained(model, peft_model, offload_folder="offload/")
32
 
33
+ model = AutoModelForCausalLM.from_pretrained(model, trust_remote_code=True, device_map="cuda:0")
34
  model = PeftModel.from_pretrained(model, peft_model)
35
 
36
+ #device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
37
+ device = 'cuda:0'
38
  model.to(device)
39
 
40
  model = model.eval()