get_device¶
- torchtune.utils.get_device(device: Optional[str] = None) device[source]¶
Function that takes an optional device string, verifies it’s correct and available given the machine and distributed settings, and returns a torch.device. If device string is not provided, this function will infer the device based on the environment.
If CUDA is available and being used, this function also sets the CUDA device.
- Parameters:
device (Optional[str]) – The name of the device to use.
- Returns:
device.
- Return type:
