torch.cuda.set_device

torch.cuda.set_device(device)

設定當前裝置。

建議不要使用此函式，而應使用 device。在大多數情況下，使用 CUDA_VISIBLE_DEVICES 環境變數更好。

引數

device (torch.device 或 int) –選定的裝置。如果此引數為負數，則此函式將不起作用。