用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上运行。