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 (裝置 或 字串) – 要設定為預設值的裝置

示例

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