指定用第几块GPU、CUDA error: invalid device ordinal、

在指定用第几块GPU的时候,在代码开头输入了

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

表示用第二块GPU,并且也只能看见第二块,然后输入:

print('# GPUs = %d' % (torch.cuda.device_count()))

这时候输出的就是1(其实一共有4块)

至此应该不会有什么问题,但是问题出就出在后面的代码:

device = 'cuda'
model.cuda()

最好用以下形式,而不是用.cuda(),因为这个无法指定第几块,它一直会用第0块

device = torch.device("cuda:3" if torch.cuda.is_available() else "cpu")
model.to(self._model_device)

 

posted @ 2021-05-10 18:46  小千北同学超爱写代码  阅读(3157)  评论(0编辑  收藏  举报