Yurii Paniv commited on
Commit
6f8eeb0
·
1 Parent(s): f0666ff

Add device

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -26,7 +26,7 @@ def translate(input_text):
26
  # iteratively generate
27
  input_text = input_text.strip()
28
  input_text = f"[INST] {input_text} [/INST]"
29
- inputs = tokenizer([input_text], return_tensors="pt")
30
 
31
  streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
32
  generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=200)
 
26
  # iteratively generate
27
  input_text = input_text.strip()
28
  input_text = f"[INST] {input_text} [/INST]"
29
+ inputs = tokenizer([input_text], return_tensors="pt").to(model.device)
30
 
31
  streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
32
  generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=200)