How to specify GPU to use

[remichu-ai]

Hi,

Just want to say this is a great library. I would like to ask if it is possible to specify the GPU to use (not using docker).

Current device is able to select between 'cuda', 'auto', etc. But is there a way to specify the GPU to use? i tried device='cuda:2' and it doesnt work.

[lckr]

Hi, does setting export CUDA_VISIBLE_DEVICES=2 to select a single gpu or export CUDA_VISIBLE_DEVICES=0,1,2 for multiple gpus help?

[remichu-ai]

It works!! Thank you very much