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

参数:

设备（设备或字符串）- 设置为默认的设备

示例：

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