Hi 
I use:

os.environ[“CUDA_VISIBLE_DEVICES”]="" to switch to cpu mode
os.environ[“CUDA_VISIBLE_DEVICES”]=“0” to cuda mode, if I have only 1 GPU
os.environ[“CUDA_VISIBLE_DEVICES”]=“0,2,5” to use only special devices (note, that in this case, pytorch will count all available devices as 0,1,2 )