Update src/app/predict.py

src/app/predict.py

@@ -3,12 +3,16 @@ import sys
 from typing import Dict
 from src.logger import logging
 from src.exception import CustomExceptionHandling
+import torch
 from transformers import pipeline
 
 
 # Load the zero-shot classification model
+device = "cuda" if torch.cuda.is_available() else "cpu"
 classifier = pipeline(
-    "zero-shot-classification", model="MoritzLaurer/ModernBERT-large-zeroshot-v2.0"
+    "zero-shot-classification",
+    model="MoritzLaurer/ModernBERT-large-zeroshot-v2.0",
+    device=device,
 )