You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by da...@apache.org on 2024/07/23 17:01:13 UTC
(beam) branch master updated: Use correct device ordinal when GPU is detected (#31951)
This is an automated email from the ASF dual-hosted git repository.
damccorm 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 9b6c805e1cd Use correct device ordinal when GPU is detected (#31951)
9b6c805e1cd is described below
commit 9b6c805e1cdb911eb11f5ade137a92b087edcf86
Author: Danny McCormick <da...@google.com>
AuthorDate: Tue Jul 23 19:01:06 2024 +0200
Use correct device ordinal when GPU is detected (#31951)
---
sdks/python/apache_beam/ml/inference/huggingface_inference.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/sdks/python/apache_beam/ml/inference/huggingface_inference.py b/sdks/python/apache_beam/ml/inference/huggingface_inference.py
index 91efcdd76a2..2934a536291 100644
--- a/sdks/python/apache_beam/ml/inference/huggingface_inference.py
+++ b/sdks/python/apache_beam/ml/inference/huggingface_inference.py
@@ -677,7 +677,7 @@ class HuggingFacePipelineModelHandler(ModelHandler[str,
self._load_pipeline_args['device'] = 'cpu'
else:
if is_gpu_available_torch():
- self._load_pipeline_args['device'] = 'cuda:1'
+ self._load_pipeline_args['device'] = 'cuda:0'
else:
_LOGGER.warning(
"HuggingFaceModelHandler specified a 'GPU' device, "