get_device
- torchtune.utils.get_device(device: Optional[str] = None) device [source]
该函数接收一个可选的设备字符串,验证它在给定机器和分布式设置下是否正确且可用,然后返回一个
device()
。如果未提供设备字符串,则该函数将根据环境推断设备。如果 CUDA 可用且正在使用,则该函数还会设置 CUDA 设备。
- 参数::
device (可选[str]) – 要使用的设备名称,例如“cuda”或“cpu”。
示例
>>> device = get_device("cuda") >>> device device(type='cuda', index=0)
- 返回::
设备
- 返回类型::