Update README.md

README.md

@@ -22,6 +22,7 @@ Use the code below to get started with the model.
 
 ```
 import math
+import torch
 import pandas as pd
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
@@ -29,6 +30,8 @@ BATCH_SIZE = 32
 ds = pd.read_csv('test.csv')
 BASE_MODEL = 'HausaNLP/afrisenti-hau-regression'
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 model = AutoModelForSequenceClassification.from_pretrained(BASE_MODEL)
 
@@ -37,7 +40,7 @@ y_preds = []
 
 for i in range(nb_batches):
   input_texts = ds[i * BATCH_SIZE: (i+1) * BATCH_SIZE]["tweet"]
-  encoded = tokenizer(input_texts, truncation=True, padding="max_length", max_length=256, return_tensors="pt").to("cuda")
+  encoded = tokenizer(input_texts, truncation=True, padding="max_length", max_length=256, return_tensors="pt").to(device)
   y_preds += model(**encoded).logits.reshape(-1).tolist()
 
 df = pd.DataFrame([ds['tweet'], ds['label'], y_preds], ["Text", "Label", "Prediction"]).T