Use correct device ordinal when GPU is detected (#31951)

apache · Jul 23, 2024 · 9b6c805 · 9b6c805
1 parent aa3cfe5
commit 9b6c805
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sdks/python/apache_beam/ml/inference/huggingface_inference.py b/sdks/python/apache_beam/ml/inference/huggingface_inference.py
@@ -677,7 +677,7 @@ def _deduplicate_device_value(self, device: Optional[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, "