ray.train.torch.get_devices#

ray.train.torch.get_devices() → List[torch.device][source]#

Gets the correct torch device list configured for this process.

Assumes that CUDA_VISIBLE_DEVICES is set and is a superset of the ray.get_gpu_ids().

Examples

Example: Launched 2 workers on the current node, each with 1 GPU

os.environ["CUDA_VISIBLE_DEVICES"] == "2,3"
ray.get_gpu_ids() == [2]
torch.cuda.is_available() == True
get_devices() == [torch.device("cuda:0")]

Example: Launched 4 workers on the current node, each with 1 GPU

os.environ["CUDA_VISIBLE_DEVICES"] == "0,1,2,3"
ray.get_gpu_ids() == [2]
torch.cuda.is_available() == True
get_devices() == [torch.device("cuda:2")]

Example: Launched 2 workers on the current node, each with 2 GPUs

os.environ["CUDA_VISIBLE_DEVICES"] == "0,1,2,3"
ray.get_gpu_ids() == [2,3]
torch.cuda.is_available() == True
get_devices() == [torch.device("cuda:2"), torch.device("cuda:3")]

PublicAPI (beta): This API is in beta and may change before becoming stable.