Update README.md

README.md

@@ -32,7 +32,7 @@ Eurus-RM-7B is trained on a mixture of [UltraInteract](https://huggingface.co/da
 ## Usage
 ```python
 from transformers import AutoTokenizer, AutoModel
-
+import torch
 
 def test(model_path):
     dataset = [ # cases in webgpt; we use the same template as Mistral-Instruct-v0.2
@@ -43,12 +43,13 @@ def test(model_path):
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
 
-    for example in dataset:
-        inputs = tokenizer(example["chosen"], return_tensors="pt")
-        chosen_reward = model(**inputs).item()
-        inputs = tokenizer(example["rejected"], return_tensors="pt")
-        rejected_reward = model(**inputs).item()
-        print(chosen_reward - rejected_reward)
+    with torch.no_grad():
+      for example in dataset:
+      inputs = tokenizer(example["chosen"], return_tensors="pt")
+      chosen_reward = model(**inputs).item()
+      inputs = tokenizer(example["rejected"], return_tensors="pt")
+      rejected_reward = model(**inputs).item()
+      print(chosen_reward - rejected_reward)
 
 test("openbmb/Eurus-RM-7b")
 # Output: 47.4404296875