Fix import and typo in device (#1)
Browse files- Fix import and typo in device (3df0acb1b8d1b5015b2882b623b7503530e95172)
Co-authored-by: rohinish <rk404@users.noreply.huggingface.co>
README.md
CHANGED
|
@@ -59,7 +59,7 @@ Additionally, `distilroberta-base` is case-sensitive model.
|
|
| 59 |
### Transformers
|
| 60 |
|
| 61 |
```python
|
| 62 |
-
from transformers import AutoTokenizer, AutoModelForSequenceClassification
|
| 63 |
import torch
|
| 64 |
|
| 65 |
tokenizer = AutoTokenizer.from_pretrained("ProtectAI/distilroberta-base-rejection-v1")
|
|
@@ -71,7 +71,7 @@ classifier = pipeline(
|
|
| 71 |
tokenizer=tokenizer,
|
| 72 |
truncation=True,
|
| 73 |
max_length=512,
|
| 74 |
-
device=torch.device("cuda" if torch.cuda.is_available() else "
|
| 75 |
)
|
| 76 |
|
| 77 |
print(classifier("Sorry, but I can't assist with that."))
|
|
|
|
| 59 |
### Transformers
|
| 60 |
|
| 61 |
```python
|
| 62 |
+
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
|
| 63 |
import torch
|
| 64 |
|
| 65 |
tokenizer = AutoTokenizer.from_pretrained("ProtectAI/distilroberta-base-rejection-v1")
|
|
|
|
| 71 |
tokenizer=tokenizer,
|
| 72 |
truncation=True,
|
| 73 |
max_length=512,
|
| 74 |
+
device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
|
| 75 |
)
|
| 76 |
|
| 77 |
print(classifier("Sorry, but I can't assist with that."))
|