get_device

torchtune.utils.get_device(device: Optional[str] = None) → device

此函数接受一个可选的设备字符串，验证其是否正确且在给定机器和分布式设置下可用，并返回一个 device()。如果未提供设备字符串，此函数将根据环境推断设备。

如果 CUDA-like 可用且正在使用，此函数还会设置 CUDA-like 设备。

参数:

device (Optional[str]) – 要使用的设备的名称，例如 “cuda”、“cpu”、“npu” 或 “xpu”。

示例

>>> device = get_device("cuda")
>>> device
device(type='cuda', index=0)

返回:

设备

返回类型:

torch.device