torch.set_default_device

torch.set_default_device(device)[源代码]

将默认的 torch.Tensor 分配到 device 上。这不会影响那些通过显式指定 device 参数调用工厂函数的情况。工厂函数调用将会像是传入了 device 作为参数一样执行。

如果只想临时更改默认设备（而非全局设置），可以使用with torch.device(device):。

默认设备最初为cpu。如果你将默认张量设备设置为其他设备（例如，cuda）而不指定设备索引，则即使调用了torch.cuda.set_device()，张量仍将分配到当前类型的默认设备上。

警告

此函数每次调用 Python 中的 torch API（不仅仅是工厂函数）时都会带来一定的性能开销。如果你遇到问题，请在 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)