Решил следующим образом: запускаю копии скрипта, на каждую видеокарту по отдельности. Перед запуском скрипта в консоль: export CUDA_VISIBLE_DEVICES="0", где 0 - номер видеокарты. Потом python script.py