torch.cuda.set_device

torch.cuda.set_device(device) [source]

Устанавливает текущий устройство.

Использование этой функции не рекомендуется в пользу device. В большинстве случаев лучше использовать переменную среды CUDA_VISIBLE_DEVICES.

Параметры:

device (torch.device или int) – выбранное устройство. Эта функция является пустой операцией, если этот аргумент отрицательный.