yeandy opened a new issue, #22712:
URL: https://github.com/apache/beam/issues/22712
### What would you like to happen?
In PytorchModelHandler's `_load_model()` function, loading a state_dict file
(`torch.load(file)`) does not map the state dict from gpu/cpu to the specified
device resulting in the following error:
```
Attempting to deserialize object on a CUDA device but
torch.cuda.is_available() is False. If you are running on a CPU-only machine,
please use torch.load with map_location=torch.device('cpu') to map your
storages to the CPU
```
Adding `map_location=device` in the handler should resolve this issue. In
addition, an automatic conversion may need to be done if there is a mismatch
between the specified device provided to the handler, and the device of the
state_dict.
We should also add a warning log as mentioned in
https://github.com/apache/beam/issues/22711.
### Issue Priority
Priority: 2
### Issue Component
Component: sdk-py-core
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
