Python运行错误：RuntimeError: CUDA error: invalid device ordinal

用pytorch跑实验需要用到cuda加速，开始默认用的是第0块gpu，后来第0块cpu被别人占用，所以只能切换到第1块gpu。
于是进行了如下修改

device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu' )

但是出现了错误：
RuntimeError: CUDA error: invalid device ordinal

经过一番查资料，作如下修改：

import os 
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu' )

如上修改过后，便可在第1块gpu上运行。