Update README.md

README.md

@@ -36,12 +36,16 @@ answers = [
 model_name = "sdadas/polish-reranker-roberta-v2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(
-    model_name, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2"
+    model_name,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2",
+    device_map="cuda"
 )
 texts = [f"{query}</s></s>{answer}" for answer in answers]
-tokens = tokenizer(texts, padding="longest", max_length=512, truncation=True, return_tensors="pt")
+tokens = tokenizer(texts, padding="longest", max_length=512, truncation=True, return_tensors="pt").to("cuda")
 output = model(**tokens)
-results = output.logits.detach().numpy()
+results = output.logits.detach().cpu().float().numpy()
 results = np.squeeze(results)
 print(results.tolist())
 ```