torch.cuda.set_device

torch.cuda.set_device(device)[源代码]

设置当前设备。

建议优先使用device而非此函数。在大多数情况下，推荐使用CUDA_VISIBLE_DEVICES环境变量。

参数

device (torch.device 或 int) – 选定的设备。如果此参数为负数，则该函数将不做任何处理。