Spaces:

jay2219
/

GemmaGuard

Runtime error

Jay Prajapati commited on Mar 23

Commit

94e805a

1 Parent(s): c83027c

fix: Default CPU Usage

Files changed (1) hide show

src/model.py CHANGED Viewed

@@ -42,7 +42,7 @@ def get_safety_score(content, safety_policy):
     correctly.
   """
-  inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
   with torch.no_grad():
     logits = model(**inputs).logits

     correctly.
   """
+  inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
   with torch.no_grad():
     logits = model(**inputs).logits