mikegarts commited on
Commit
29e5c9b
1 Parent(s): 2c32fbf

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -17,7 +17,7 @@ def generate(prompt):
17
  model, tokenizer = get_model()
18
 
19
  input_context = prompt
20
- input_ids = tokenizer.encode(input_context, return_tensors="pt").to('cpu')
21
 
22
  outputs = model.generate(
23
  input_ids=input_ids,
 
17
  model, tokenizer = get_model()
18
 
19
  input_context = prompt
20
+ input_ids = tokenizer.encode(input_context, return_tensors="pt").to(model.device)
21
 
22
  outputs = model.generate(
23
  input_ids=input_ids,