yeandy commented on code in PR #22795:
URL: https://github.com/apache/beam/pull/22795#discussion_r954205790
##########
sdks/python/apache_beam/ml/inference/pytorch_inference.py:
##########
@@ -40,11 +41,32 @@
def _load_model(
model_class: torch.nn.Module, state_dict_path, device, **model_params):
model = model_class(**model_params)
- model.to(device)
+
+ if device == torch.device('cuda') and not torch.cuda.is_available():
+ logging.warning(
+ "Model handler specified a 'GPU' device, but GPUs are not available. "
\
+ "Switching to CPU.")
+ device = torch.device('cpu')
+
file = FileSystems.open(state_dict_path, 'rb')
- model.load_state_dict(torch.load(file))
+ try:
+ logging.info(
+ "Loading state_dict_path %s onto a %s device", state_dict_path, device)
+ state_dict = torch.load(file, map_location=device)
+ except RuntimeError as e:
+ message = "Loading the model onto a GPU device failed due to an " \
Review Comment:
Well, I can't say that the error message will always be: `Attempting to
deserialize object on a CUDA device but torch.cuda.is_available() is False`. Be
it will definitely stem from the `torch.load(file, map_location=device)` call.
So that should provide enough context as to where the exception has occured.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@beam.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
