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)