Update src/modernbert/classifier.py

src/modernbert/classifier.py

@@ -15,7 +15,6 @@ classifier = pipeline(
     "zero-shot-classification",
     model="MoritzLaurer/ModernBERT-large-zeroshot-v2.0",
     torch_dtype=torch.bfloat16,
-    device="cpu"
 )
 
 
@@ -23,7 +22,7 @@ def ZeroShotTextClassification(
     text_input: str, candidate_labels: str, multi_label: bool
 ) -> Dict[str, float]:
     """
-    Performs zero-shot classification on the given text input.
+    Performs zero-shot classification on the given text input and candidate labels.
 
     Args:
         - text_input: The input text to classify.