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