Add Generic to PredictionResult

apache · Jun 10, 2022 · a273b05 · a273b05
1 parent 52581a6
commit a273b05
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py
@@ -64,8 +64,8 @@
 
 @dataclass
 class PredictionResult:
-  example: _INPUT_TYPE
-  inference: _OUTPUT_TYPE
+  example: Generic[_INPUT_TYPE]
+  inference: Generic[_OUTPUT_TYPE]
 
 
 def _to_milliseconds(time_ns: int) -> int: