Update README.md

README.md

@@ -15,14 +15,14 @@ Note: 1. Remember to use the formulation of Anthropic/hh-rlhf dataset for infere
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-rm_tokenizer = AutoTokenizer.from_pretrained(rm_tokenizer_path)
+rm_tokenizer = AutoTokenizer.from_pretrained('Ray2333/gpt2-large-harmless-reward_model')
 reward_model = AutoModelForSequenceClassification.from_pretrained(
-                reward_peft_path1,
+                'Ray2333/gpt2-large-harmless-reward_model',
                 num_labels=1, torch_dtype=torch.bfloat16,
-                device_map=gpu_id1,
+                device_map=0,
                 )
 q, a = "\n\nHuman: I just came out of from jail, any suggestion of my future? \n\nAssistant:", "Go back to jail you scum"
 inputs = rm_tokenizer(q, a, return_tensors='pt', truncation=True)
 with torch.no_grad():
-  reward = reward_model(**(inputs.to(gpu_id1))).logits[0].cpu().detach().item()
+  reward = reward_model(**(inputs.to(0))).logits[0].cpu().detach().item()
 ```