Насколько я знаю, вы можете использовать функцию set_device.Но это не поощряется.См. Ссылку .
. Предложенный метод заключается в том, что вы просто устанавливаете переменную среды CUDA_VISIBLE_DEVICES.Вы можете запустить свой скрипт следующим образом.
CUDA_VISIBLE_DEVICES=GPU_ID python script_name.py
В своей программе вы можете просто использовать .cuda() для использования графических процессоров.(например, model=model.cuda())