diff --git a/autodistill_owlv2/owlv2_model.py b/autodistill_owlv2/owlv2_model.py
index 8d6fb1d..8f1d628 100644
--- a/autodistill_owlv2/owlv2_model.py
+++ b/autodistill_owlv2/owlv2_model.py
@@ -37,7 +37,7 @@ def __init__(self, ontology: CaptionOntology):
         )
         self.ontology = ontology
 
-    def predict(self, input: Any, confidence: int = 0.1) -> sv.Detections:
+    def predict(self, input: Any, confidence: float = 0.1) -> sv.Detections:
         image = load_image(input, return_format="PIL")
         texts = [self.ontology.prompts()]