Update webdemo.py

Clouditera · Apr 25, 2024 · 80a8b69 · 80a8b69
1 parent 62270d1
commit 80a8b69
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/webdemo/webdemo.py b/webdemo/webdemo.py
@@ -105,7 +105,7 @@ def evaluate(
             return
         prompt = reformat_sft(instruction, "")
 
-        inputs = tokenizer(prompt, return_tensors="pt")
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
         input_ids = inputs["input_ids"]
 
         if not (1 > temperature > 0):