CCCCCC commited on
Commit
19d9cd8
1 Parent(s): 892aa5d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -10,9 +10,10 @@ model_path = 'THUDM/BPO'
10
 
11
  device = 'cuda:0'
12
 
13
- tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, add_prefix_space=True)
14
- model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, device_map=device, load_in_8bit=True)
15
- model = model.eval()
 
16
 
17
 
18
  DESCRIPTION = """This Space demonstrates model [BPO](https://huggingface.co/THUDM/BPO), which is built on LLaMA-2-7b-chat.
 
10
 
11
  device = 'cuda:0'
12
 
13
+ if torch.cuda.is_available():
14
+ tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, add_prefix_space=True)
15
+ model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, device_map=device, load_in_8bit=True)
16
+ model = model.eval()
17
 
18
 
19
  DESCRIPTION = """This Space demonstrates model [BPO](https://huggingface.co/THUDM/BPO), which is built on LLaMA-2-7b-chat.