Trayvon commited on
Commit
c81ffe4
1 Parent(s): d98d451

Update src/demo.py

Browse files
Files changed (1) hide show
  1. src/demo.py +1 -1
src/demo.py CHANGED
@@ -40,7 +40,7 @@ def generate(input_text, sys_prompt, temperature, max_new_tokens) -> str:
40
  '''
41
  input_str = sys_prompt + input_text + " [/INST]"
42
 
43
- input_ids = tokenizer(input_str, return_tensors="pt")
44
 
45
  streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
46
 
 
40
  '''
41
  input_str = sys_prompt + input_text + " [/INST]"
42
 
43
+ input_ids = tokenizer(input_str, return_tensors="pt").to("cuda")
44
 
45
  streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
46