Привязать Torch Transformers Classifier Model к cudа с определенным индексом

DeepPavlov Library Models
1

Привет!

Я хочу воспользоваться Torch Transformers Classifier Model на сервере с несколькими видеокартами, однако при передаче аргумента ‘gpu’ в ‘device’ модель помещается на базовую карточку. Полагаю причиной этой является 75 строчка в torch_model.py. Есть ли способы как-то это обойти и разместить модель на определнной вдиелокарте?

Я со своей стороны пробовал после инициализации модели просто поменять значение параметра device у объекта классифиатора, это получалось сделать, однако при запуске обучения выдается ошибка RuntimeError: arguments are located on different GPUs

2

Привет! сорри, что не сразу ответили.

Можно - и вы можете прямо сейчас сделать PR к нам, чтобы это попало в библиотеку :wink:

3

Видимые карточки можно управлять переменной окружения CUDA_VISIBLE_DEVICES.
Например, можно разрешить запуск только на одной карте номер 2 (номер из nvidia-smi):

export CUDA_VISIBLE_DEVICES=2; python my_script.py