jslin09 commited on
Commit
34ea07c
·
verified ·
1 Parent(s): e026ea4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -6,7 +6,7 @@ from transformers import pipeline, set_seed
6
  import random
7
  # import torch
8
 
9
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
10
  model_name = "bloom-560m"
11
  model = AutoModelForCausalLM.from_pretrained(f'jslin09/{model_name}-finetuned-fraud').to('cuda')
12
  tokenizer = BloomTokenizerFast.from_pretrained(f'bigscience/{model_name}', bos_token = '<s>', eos_token = '</s>', pad_token = '<pad>')
 
6
  import random
7
  # import torch
8
 
9
+ # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
10
  model_name = "bloom-560m"
11
  model = AutoModelForCausalLM.from_pretrained(f'jslin09/{model_name}-finetuned-fraud').to('cuda')
12
  tokenizer = BloomTokenizerFast.from_pretrained(f'bigscience/{model_name}', bos_token = '<s>', eos_token = '</s>', pad_token = '<pad>')