torch.cuda.set_device

torch.cuda.set_device(device)

设置当前设备。

不建议使用此函数，推荐使用 device。在大多数情况下，使用 CUDA_VISIBLE_DEVICES 环境变量更好。

参数:

device (torch.device 或 int) – 所选设备。如果此参数为负数，则此函数不执行任何操作。