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 呼叫（不僅僅是工廠函式）帶來輕微的效能開銷。如果這對您造成了問題，請在 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)