protectai
/

distilroberta-base-rejection-v1

Text Classification

Generated from Trainer

Inference Endpoints

Model card Files Files and versions Community

Fix import and typo in device

#1

by rk404 - opened Mar 11

base: refs/heads/main

←

from: refs/pr/1

Discussion Files changed

Files changed (1) hide show

README.md +2 -2

README.md CHANGED Viewed

@@ -59,7 +59,7 @@ Additionally, `distilroberta-base` is case-sensitive model.
 ### Transformers
 ```python
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 tokenizer = AutoTokenizer.from_pretrained("ProtectAI/distilroberta-base-rejection-v1")
@@ -71,7 +71,7 @@ classifier = pipeline(
   tokenizer=tokenizer,
   truncation=True,
   max_length=512,
-  device=torch.device("cuda" if torch.cuda.is_available() else "CPU"),
 )
 print(classifier("Sorry, but I can't assist with that."))

 ### Transformers
 ```python
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 import torch
 tokenizer = AutoTokenizer.from_pretrained("ProtectAI/distilroberta-base-rejection-v1")
   tokenizer=tokenizer,
   truncation=True,
   max_length=512,
+  device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
 )
 print(classifier("Sorry, but I can't assist with that."))