Is there a way to choose which GPU a model uses?

[dominic-simon]

I'm trying to figure out if there is a way to set which GPU is used (ex. 'cuda:0' or 'cuda:1') during inference when using the Pytorch Faster-RCNN Object Detector. There is a parameter, device_type, that can be set to either 'cpu' or 'gpu'. When choosing 'gpu', however, there is no way to choose which GPU is used. It just uses the GPU with the lowest ID (i.e. 0). Is there a way to set the specific GPU that the model will run on?

I've tried setting device_type to 'cuda:1' and I've tried using .to('cuda:1'), but neither work. Any suggestion would be appreciated.

[beat-buesser]

Hi @dominic-simon Sorry for the delayed response. I think that's a good question, I don't think this has been tested so far. The current approach is based on using the CUDA environment variable CUDA_VISIBLE_DEVICES and set it to define which GPUs are visible. Can you try setting export CUDA_VISIBLE_DEVICES=1 and setting device to gpu before running your code?