gpu

tasks/text.py

@@ -2,7 +2,7 @@ from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
-import random
+import torch
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
@@ -13,6 +13,7 @@ classifier = pipeline(
     task="text-classification",
     model="axel-darmouni/test_modernbert_finetuned",
     tokenizer="answerdotai/ModernBERT-base",
+    device="cuda:0" if torch.cuda.is_available() else "cpu"
 )