torchtnt.utils.device.get_device_from_env

torchtnt.utils.device.get_device_from_env() → device

Function that gets the torch.device based on the current environment.

This currently supports only CPU and GPU devices. If CUDA is available, this function also sets the CUDA device.

Within a distributed context, this function relies on the LOCAL_RANK environment variable to be made available by the program launcher for setting the appropriate device index.

Raises:RuntimeError – If LOCAL_RANK is outside the range of available GPU devices.