This is an automated email from the ASF dual-hosted git repository.
bhulette pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new eedff4b996c Change return type for PytorchInferenceRunner (#17460)
eedff4b996c is described below
commit eedff4b996c042d4e5a61e42a2da5f776ff42afd
Author: Anand Inguva <[email protected]>
AuthorDate: Mon Apr 25 23:09:22 2022 +0000
Change return type for PytorchInferenceRunner (#17460)
---
sdks/python/apache_beam/ml/inference/pytorch.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/sdks/python/apache_beam/ml/inference/pytorch.py
b/sdks/python/apache_beam/ml/inference/pytorch.py
index 62741a899e8..438582d73dc 100644
--- a/sdks/python/apache_beam/ml/inference/pytorch.py
+++ b/sdks/python/apache_beam/ml/inference/pytorch.py
@@ -37,7 +37,7 @@ class PytorchInferenceRunner(InferenceRunner):
self._device = device
def run_inference(self, batch: List[torch.Tensor],
- model: torch.nn.Module) -> Iterable[torch.Tensor]:
+ model: torch.nn.Module) -> Iterable[PredictionResult]:
"""
Runs inferences on a batch of Tensors and returns an Iterable of
Tensor Predictions.
