How to change the default device of GPU? device_ids[0]

10

I have tried to set CUDA_VISIBLE_DEVICES in shell, then I run a simple script to test if the setting has taken effect, unfortunately, it seems does not work.

CUDA_VISIBLE_DEVICES=3; python test.py

the script test.py is

import torch
print(torch.cuda.current_device())

the above script still shows that current device is 0.

I find that torch.cuda.set_device(device_num) works fine in setting the desired GPU to use.

16 Likes