Spaces:
Sleeping
Sleeping
Update src/demo.py
Browse files- 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 |
|