快捷方式

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)
返回::

设备

返回类型::

torch.device

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源