Update README.md

README.md

@@ -37,11 +37,12 @@ We evaluate GRM 2B on the [reward model benchmark](https://huggingface.co/spaces
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
+device = 'cuda:2'
 # load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained('Ray2333/GRM-Gemma-2B-sftreg')
 reward_model = AutoModelForSequenceClassification.from_pretrained(
                 'Ray2333/GRM-Gemma-2B-sftreg', torch_dtype=torch.float16,  trust_remote_code=True, 
-                device_map=0,
+                device_map=device,
                 )
 message = [
   {'role': 'user', 'content': "I'm going to go out to a movie, but I need someone to chat with my daughter and pretend to be me while she's home alone.  But I can't do that while I'm at the movie.  Can you help by impersonating me by chat with her?"},
@@ -54,7 +55,7 @@ kwargs = {"padding": 'max_length', "truncation": True, "return_tensors": "pt"}
 tokens = tokenizer.encode_plus(message_template, **kwargs)
 
 with torch.no_grad():
-  _, _, reward_tensor = model(tokens["input_ids"][0].to(model.device), attention_mask=tokens["attention_mask"][0].to(model.device)).logits.reshape(-1)
+  _, _, reward_tensor = reward_model(tokens["input_ids"][0].view(1,-1).to(device), attention_mask=tokens["attention_mask"][0].view(1,-1).to(device))
   reward = reward_tensor.cpu().detach().item()
 ```