dar-tau commited on
Commit
90e0ba3
·
verified ·
1 Parent(s): c01c7c6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -46,9 +46,10 @@ Assistant: girlfriend;mother;father;friend
46
 
47
  # setup
48
  torch.set_grad_enabled(False)
 
49
  model_name = "TheBloke/OpenHermes-2.5-Mistral-7B-GPTQ"
50
  token = os.environ['hf_token']
51
- pipe = pipeline("text-generation", model=model_name, device="cuda")
52
  pipe.tokenizer.chat_template = chatml_template # TheBloke says this is the right template for this model
53
  generate_kwargs = {'max_new_tokens': 20}
54
 
 
46
 
47
  # setup
48
  torch.set_grad_enabled(False)
49
+ device = "cpu"
50
  model_name = "TheBloke/OpenHermes-2.5-Mistral-7B-GPTQ"
51
  token = os.environ['hf_token']
52
+ pipe = pipeline("text-generation", model=model_name, device=device)
53
  pipe.tokenizer.chat_template = chatml_template # TheBloke says this is the right template for this model
54
  generate_kwargs = {'max_new_tokens': 20}
55