to cuda

2025-03-22 08:15:33 +08:00 · 2025-03-22 08:15:33 +08:00 · 876f1f37a1
commit 876f1f37a1
parent 679f1bd941
1 changed files with 1 additions and 0 deletions
--- a/article_analyzer.py
+++ b/article_analyzer.py
@ -26,6 +26,7 @@ class ArticleClassification(BaseModel):
 class ArticleAnalyzer:
    def __init__(self):
        self.model = AutoModelForSequenceClassification.from_pretrained("ProsusAI/finbert", num_labels=3, cache_dir=None)
+        self.model.to("cuda")
        self.base_prompt = """
            Classify the following article into one of these categories: 
            - Regulatory News