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
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, e.g. “cuda” or “cpu”.
Example
>>> device = get_device("cuda") >>> device device(type='cuda', index=0)
- Returns:
Device
- Return type:
