Update app.py

app.py

@@ -15,7 +15,7 @@ tokenizer = BloomTokenizerFast.from_pretrained(f'bigscience/{model_name}', bos_t
 def rnd_generate(prompt):
     rnd_seed = random.randint(10, 500)
     set_seed(rnd_seed)
-    inputs = tokenizer(prompt, return_tensors="pt") # 回傳的張量使用 Pytorch的格式。如果是 Tensorflow 格式的話，則指定為 "tf"。
+    inputs = tokenizer(prompt, return_tensors="pt").to('cuda') # 回傳的張量使用 Pytorch的格式。如果是 Tensorflow 格式的話，則指定為 "tf"。
     results = model.generate(inputs["input_ids"],
                        max_length=500,
                        num_return_sequences=1, # 產生 1 個句子回來。
@@ -28,7 +28,7 @@ def rnd_generate(prompt):
 @spaces.GPU
 def generate(prompt):
     result_length = len(prompt) + 4
-    inputs = tokenizer(prompt, return_tensors="pt") # 回傳的張量使用 Pytorch的格式。如果是 Tensorflow 格式的話，則指定為 "tf"。
+    inputs = tokenizer(prompt, return_tensors="pt").to('cuda') # 回傳的張量使用 Pytorch的格式。如果是 Tensorflow 格式的話，則指定為 "tf"。
     results = model.generate(inputs["input_ids"],
                        num_return_sequences=2, # 產生 2 個句子回來。
                        max_length=result_length,