Closed amogkam closed 1 year ago
GPU device ids can be specified with an integer index, but may also be specified as strings.
This PR ensures that both cases are supported by root_device. The code is taken from what is being done in Ray Train: https://sourcegraph.com/github.com/ray-project/ray/-/blob/python/ray/train/torch/train_loop_utils.py?L470-498
root_device
Closes https://github.com/ray-project/ray_lightning/issues/236
GPU device ids can be specified with an integer index, but may also be specified as strings.
This PR ensures that both cases are supported by
root_device
. The code is taken from what is being done in Ray Train: https://sourcegraph.com/github.com/ray-project/ray/-/blob/python/ray/train/torch/train_loop_utils.py?L470-498Closes https://github.com/ray-project/ray_lightning/issues/236