Pytorch Cheat Sheet

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....

October 16, 2022 · 1 min · SY Chou