快捷方式

torch.set_default_device

torch.set_default_device(device)[源代码]

将默认的 torch.Tensor 分配到 device 上。这不会影响带有显式 device 参数的工厂函数调用。工厂调用将执行,就好像它们传递了 device 作为参数一样。

要临时更改默认设备而不是全局设置它,请改用 with torch.device(device):

默认设备最初为 cpu。如果您将默认张量设备设置为另一个设备(例如 cuda)且没有设备索引,则即使在调用 torch.cuda.set_device() 之后,张量仍将在该设备类型的当前设备上分配。

警告

此函数会对每次对 torch API 的 Python 调用(不仅仅是工厂函数)造成轻微的性能损失。如果这给您带来问题,请在 https://github.com/pytorch/pytorch/issues/92701 上发表评论

注意

这不会影响创建与输入共享相同内存的张量的函数,例如:torch.from_numpy()torch.frombuffer()

参数

device (device字符串) – 要设置为默认的设备

示例

>>> torch.get_default_device()
device(type='cpu')
>>> torch.set_default_device('cuda')  # current device is 0
>>> torch.get_default_device()
device(type='cuda', index=0)
>>> torch.set_default_device('cuda')
>>> torch.cuda.set_device('cuda:1')  # current device is 1
>>> torch.get_default_device()
device(type='cuda', index=1)
>>> torch.set_default_device('cuda:1')
>>> torch.get_default_device()
device(type='cuda', index=1)

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取初学者和高级开发人员的深入教程

查看教程

资源

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

查看资源