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（不仅限于工厂函数）施加轻微的性能开销。如果这给您带来了问题，请在 pytorch/pytorch#92701 上发表评论。

注意

这不会影响创建与输入共享内存的张量的函数，例如：torch.from_numpy() 和 torch.frombuffer()

参数:

device (device 或 string) – 要设置为默认值的设备

示例

>>> 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)