StevenChen16 commited on
Commit
61454b0
1 Parent(s): 52d083e

update app.py to use cuda

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -4,7 +4,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
4
  # 加载模型和tokenizer
5
  model_name = "StevenChen16/llama3-8b-Lawyer"
6
  tokenizer = AutoTokenizer.from_pretrained(model_name)
7
- model = AutoModelForCausalLM.from_pretrained(model_name)
8
 
9
  def generate_response(input_text):
10
  inputs = tokenizer(input_text, return_tensors="pt")
 
4
  # 加载模型和tokenizer
5
  model_name = "StevenChen16/llama3-8b-Lawyer"
6
  tokenizer = AutoTokenizer.from_pretrained(model_name)
7
+ model = AutoModelForCausalLM.from_pretrained(model_name).to('cuda')
8
 
9
  def generate_response(input_text):
10
  inputs = tokenizer(input_text, return_tensors="pt")