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)