pylints

apache · Jul 24, 2023 · 4bdab80 · 4bdab80
1 parent 2f95adc
commit 4bdab80
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/sdks/python/apache_beam/ml/inference/huggingface_inference.py b/sdks/python/apache_beam/ml/inference/huggingface_inference.py
@@ -33,7 +33,6 @@
 from apache_beam.ml.inference import utils
 from apache_beam.ml.inference.base import ModelHandler
 from apache_beam.ml.inference.base import PredictionResult
-from apache_beam.ml.inference.base import PredictionT
 from apache_beam.ml.inference.pytorch_inference import _convert_to_device
 from transformers import AutoModel
 from transformers import TFAutoModel
@@ -166,7 +165,7 @@ def __init__(
       framework: str,
       device: str = "CPU",
       *,
-      inference_fn: Optional[Callable[..., Iterable[PredictionT]]] = None,
+      inference_fn: Optional[Callable[..., Iterable[PredictionResult]]] = None,
       load_model_args: Optional[Dict[str, Any]] = None,
       inference_args: Optional[Dict[str, Any]] = None,
       min_batch_size: Optional[int] = None,
@@ -345,7 +344,7 @@ def __init__(
       model_class: Union[AutoModel, TFAutoModel],
       device: str = "CPU",
       *,
-      inference_fn: Optional[Callable[..., Iterable[PredictionT]]] = None,
+      inference_fn: Optional[Callable[..., Iterable[PredictionResult]]] = None,
       load_model_args: Optional[Dict[str, Any]] = None,
       inference_args: Optional[Dict[str, Any]] = None,
       min_batch_size: Optional[int] = None,