how to select specific CUDA

[xzitlou]

Hello there, I'm trying to change the CUDA selection, I saw this:

def run_inference(args, gpu_num, gpu_no, **kwargs):
    ...
    model = model.cuda(gpu_no)
    ...

if __name__ == '__main__':
    ...
    rank, gpu_num = 0, 1
    run_inference(args, gpu_num, rank)

I tried to change rank to 3 so I can use my third CUDA and also adding gpu_no=3 as param:

Traceback (most recent call last):
  File "scripts/evaluation/inference.py", line 138, in <module>
    run_inference(args, gpu_num, rank, gpu_no=3)
TypeError: run_inference() got multiple values for argument 'gpu_no'

[RuojiWang]

I happened to encounter this problem. After reading code and testing, I did not need to modify the original inference code. I just add a line before the run_inference method in the main function: "os.environ['CUDA_VISIBLE_DEVICES'] = '7'" , so that pytorch will run on the graphics card numbered 7, which is the eighth graphics card.

In this way, I can choose the graphics card to use without setting parameters such as rank, gpu_num and so on.

The following is the sample code for inference.py. I only added the line os.environ['CUDA_VISIBLE_DEVICES'] = '7'.