torch.cuda.set_device

torch.cuda.set_device(device)

设置当前设备。

使用此功能被建议，以 device 代替。在大多数情况下，使用 CUDA_VISIBLE_DEVICES 环境变量会更好。

参数:

设备（torch.device 或 int）- 选择设备。如果此参数为负，则此函数为空操作。