get_device¶
- torchtune.utils.get_device(device: Optional[str] = None) device[source]¶
Function that takes or device or device string, verifies it’s correct and availabe given the machine and distributed settings, and returns a torch.device.
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
