vishesh-t27 commited on
Commit
e77ee07
·
verified ·
1 Parent(s): 3972b10

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +3 -2
README.md CHANGED
@@ -27,13 +27,14 @@ import json
27
 
28
  model_name = "FrontiersMind/Nandi-Mini-150M-GuardRails"
29
 
30
- tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
31
 
 
32
  model = AutoModelForCausalLM.from_pretrained(
33
  model_name,
34
  trust_remote_code=True,
35
  dtype=torch.bfloat16
36
- ).cuda().eval()
37
 
38
 
39
  def classify_safety(prompt, response=None):
 
27
 
28
  model_name = "FrontiersMind/Nandi-Mini-150M-GuardRails"
29
 
30
+ device = "cuda" if torch.cuda.is_available() else "cpu"
31
 
32
+ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
33
  model = AutoModelForCausalLM.from_pretrained(
34
  model_name,
35
  trust_remote_code=True,
36
  dtype=torch.bfloat16
37
+ ).to(device).eval()
38
 
39
 
40
  def classify_safety(prompt, response=None):