torch.cuda.set_device

torch.cuda.set_device(device) [исходный код]

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

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

Параметры:

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