type fix

terrier-org · Sep 18, 2024 · f544d2f · f544d2f
1 parent 9133a9a
commit f544d2f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pyterrier/apply_base.py b/pyterrier/apply_base.py
@@ -294,7 +294,7 @@ def _feature_fn(row):
             pipe = pt.terrier.Retriever(index) >> pt.apply.doc_features(_feature_fn) >> pt.LTRpipeline(xgBoost())
     """
     def __init__(self,
-        fn: Callable[[pd.Series, pt.model.IterDictRecord], np.array],
+        fn: Callable[[Union[pd.Series, pt.model.IterDictRecord]], np.array],
         *,
         verbose: bool = False
     ):