get_device

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

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

如果 CUDA 或类似设备可用且正在使用，此函数也会设置相应的 CUDA 或类似设备。

参数:

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

示例

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

返回:

设备

返回类型:

torch.device