pytorch显卡号指定

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文件中试了不行！！！

请路过的大佬，指点迷津！！多谢！！