From a273b052c759e81df8a06b59bc17e15123293933 Mon Sep 17 00:00:00 2001
From: Andy Ye <andyye333@gmail.com>
Date: Fri, 10 Jun 2022 12:34:41 -0400
Subject: [PATCH] Add Generic to PredictionResult

---
 sdks/python/apache_beam/ml/inference/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py
index 031d58b47e2e..4eb2ee05d483 100644
--- 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: