zhengr commited on
Commit
15dcaf2
·
verified ·
1 Parent(s): 0ae1c3d

Update app.py

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