Set Up Visible Devices For PyTorch

It’s well-known that we can use os.environ["CUDA_VISIBLE_DEVICES"]="1,2" to determine which GPU can be used by the program, but as for PyTorch, most of answers says that there is no way to set visible devices in the python code for PyTorch.

However, I found os.environ.setdefault can do this.

import os
import torch

gpus = [1, 2]
os.environ.setdefault("CUDA_VISIBLE_DEVICES", ','.join(map(str, gpus)))
print(f"PyTorch detected number of availabel devices: {torch.cuda.device_count()}")

The above code will show

PyTorch detected number of availabel devices: 2