Update README.md

README.md

@@ -166,14 +166,14 @@ def build_inputs_prompt(current_Dialogue):
 
 path = "HeAAAAA/RoleRM"
 
-model = AutoModelForCausalLM.from_pretrained(path)
+model = AutoModelForCausalLM.from_pretrained(path).to("cuda")
 tokenizer = AutoTokenizer.from_pretrained(path)
 
 
 inputs_prompt1 = build_inputs_prompt(current_Dialogue1) 
 inputs_prompt2 = build_inputs_prompt(current_Dialogue2) 
 
-inputs = tokenizer(inputs_prompt1, return_tensors="pt").to("cuda:2")
+inputs = tokenizer(inputs_prompt1, return_tensors="pt")
 outputs = model.generate(
     **inputs,
     max_new_tokens=200,
@@ -190,7 +190,7 @@ new_text = tokenizer.decode(new_tokens, skip_special_tokens=True)
 
 print(new_text)
 
-inputs = tokenizer(inputs_prompt2, return_tensors="pt").to("cuda:2")
+inputs = tokenizer(inputs_prompt2, return_tensors="pt").to("cuda")
 outputs = model.generate(
     **inputs,
     max_new_tokens=200,