pytorch显卡号默认选择0
1.torch.cuda.set_device(1)
此法适合在生成网络对象之前使用,即放在最顶层的文件中,而且要import torch
2.self.net_bone = self.net_bone.cuda(i)
初始化模型时使用,但是要注意将数据存放的GPU也要修改,不然可能会报错,一般在程序中寻找sal_image, sal_label, sal_edge = sal_image.cuda(), sal_label.cuda(), sal_edge.cuda()这些字眼,修改数据的显卡号,即可
RuntimeError: Expected tensor for argument #1 ‘input’ to have the same device as tensor for argument #2 ‘weight’; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)
上述方法亲测有效
×× export CUDA_VISIBLE_DEVICES=0 #这里是要使用的GPU编号,正常的话是从0开始
这个方法在.sh文件中试了不行!!!