SengTak commited on
Commit
1fb0958
·
verified ·
1 Parent(s): 1f7a1ec

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -11,7 +11,7 @@ best_clf = joblib.load("./saved/classifier_llama32.joblib")
11
 
12
  encoder = LLMEncoder.from_pretrained(
13
  "SengTak/lec_harmfulprompclassifierencoder",
14
- device_map="cpu",
15
  torch_dtype=torch.bfloat16,
16
  #torch_dtype=torch,
17
  #cache_dir=cache_dir
 
11
 
12
  encoder = LLMEncoder.from_pretrained(
13
  "SengTak/lec_harmfulprompclassifierencoder",
14
+ device_map="cuda" if torch.cuda.is_available() else "cpu",
15
  torch_dtype=torch.bfloat16,
16
  #torch_dtype=torch,
17
  #cache_dir=cache_dir