Update app.py

app.py

@@ -4,11 +4,11 @@ from transformers import AutoModelForCausalLM
 from transformers import BloomTokenizerFast
 from transformers import pipeline, set_seed
 import random
-import torch
+# import torch
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model_name = "bloom-560m"
-model = AutoModelForCausalLM.from_pretrained(f'jslin09/{model_name}-finetuned-fraud').to(device)
+model = AutoModelForCausalLM.from_pretrained(f'jslin09/{model_name}-finetuned-fraud').to('cuda')
 tokenizer = BloomTokenizerFast.from_pretrained(f'bigscience/{model_name}', bos_token = '<s>', eos_token = '</s>', pad_token = '<pad>')
 
 def rnd_generate(prompt):